'Reward function that rewards the model for playing within a given key.
Any note within the key is given equal reward, which can cause the model to
learn random sounding compositions.
Args:
action: One-hot encoding of the chosen action.
key: The numeric values of notes belonging to this key. Defaults to C
Major if not provided.
Returns:
Float reward value.'

def reward_key_distribute_prob(self, action, key=None):

    if (key is None):
        key = rl_tuner_ops.C_MAJOR_KEY
    reward = 0
    action_note = np.argmax(action)
    if (action_note in key):
        num_notes_in_key = len(key)
        extra_prob = (1.0 / num_notes_in_key)
        reward = extra_prob
    return reward
