'Reward function that trains the model to play a scale.
Gives rewards for increasing notes, notes within the desired scale, and two
consecutive notes from the scale.
Args:
obs: A one-hot encoding of the observed note.
action: A one-hot encoding of the chosen action.
scale: The scale the model should learn. Defaults to C Major if not
provided.
Returns:
Float reward value.'

def reward_scale(self, obs, action, scale=None):

    if (scale is None):
        scale = rl_tuner_ops.C_MAJOR_SCALE
    obs = np.argmax(obs)
    action = np.argmax(action)
    reward = 0
    if (action == 1):
        reward += 0.1
    if ((action > obs) and (action < (obs + 3))):
        reward += 0.05
    if (action in scale):
        reward += 0.01
        if (obs in scale):
            action_pos = scale.index(action)
            obs_pos = scale.index(obs)
            if ((obs_pos == (len(scale) - 1)) and (action_pos == 0)):
                reward += 0.8
            elif (action_pos == (obs_pos + 1)):
                reward += 0.8
    return reward
