Reward transformations

[jbrea]

Sometimes one wants to give a transformed reward to the learner, but keep the true reward given by the environment for evaluation purposes. For example Dopamine clamps all rewards to [-1, 1] and I believe some of our methods are unstable in the Atari domain, because we don't clip the rewards. Where would it be best to transform rewards? Should we add a POST_OBSERVE hook, or allow for applying the transformation when observations are put into buffers or just before the actual learning takes place?

[findmyway]

I think the best place to do reward transform is in the preprocessor of a wrapped env:

https://github.com/JuliaReinforcementLearning/ReinforcementLearningZoo.jl/blob/master/src/experiments/atari.jl#L27-L30

We can add a preprocessor like this (this part should be added in RLBase:

struct RewardOverriddenObs{O, R}
    obs::O
    reward::R
end

@forward RewardOverriddenObs.obs ActionStyle,
get_legal_actions,
get_legal_actions_mask,
get_terminal,
get_state

RLBase.get_reward(obs::RewardOverriddenObs) = obs.reward

clip_reward(obs) = RewardOverriddenObs(obs, clamp(get_reward(obs), -1, 1)

To record the real reward, we need to write another custom hook extend the TotalRewardPerEpisode (Edit: I'm not sure whether to make it as the default behavior, need decision):

function (hook::TotalRewardPerEpisode)(::PostActStage, agent, env, obs::RewardOverriddenObs)
    hook.reward += get_reward(obs.obs)
end

Thanks for reminding this. I thought it was just a trick used in NatureDQN. I'll give it a try and see if it will improve the results further.

By the way, in some implementations, sign is used, do you know if there's any significant difference on the result? See the one in rlpyt

[jbrea]

We can add a preprocessor

Sounds good!

I thought it was just a trick

I think it allows to use the same learning rate across different games. E.g. in pong the immediate rewards are -1, 0 or 1, whereas in other games the immediate rewards can be more than 1000. I was running IQN on a random game (time pilot) and saw it becoming unstable.

By the way, in some implementations, sign is used, do you know if there's any significant difference on the result?

In environments where the absolute value of the reward is always larger than 1, there is obviously no difference 😄. Whether there are atari games with smaller rewards I don't know.

Btw. not surprisingly, this reward clipping is pretty crude and can strongly affect the results.

[findmyway]

Thanks. I'll make a PR soon.

[findmyway]

The result of Dopamine_Rainbow_Atari(time_pilot) now:

[jbrea]

Cool!!! I guess this is roughly comparable to dopamine. I guess their 200 iterations correspond to our 50_000_000 steps. But there are many things to make it difficult to compare on Atari; for reference: https://arxiv.org/abs/1908.04683 and https://github.com/valeoai/rainbow-iqn-apex.