Why train value network on cumulative discounted return?

Question

I recently read a lot of code regarding policy gradients, in particular, PPO and TRPO. To reduce variance the common approach is to train a network function as a baseline (value function, actor-critic style). Usually, this is done by training the value network on the discounted cumulative reward/return. Reference in tensorflow agents, for instance: https://github.com/tensorflow/agents/blob/master/agents/algorithms/ppo/ppo.py#L437-L446

The thing I find a little odd is that the value network usually does not know the current timestep. It is only given the current state as input. But using discounted cumulative return the value is much higher in an early timestep than in a later timestep. So why would you want to train the value network on the cumulative discounted return? Imagine having a state occurring at the start of an episode and close to the end, the cumulative discounted reward would be very different. Am I missing something here?

Thanks, Magnus

Answer 1

Formally, in a fully observed markov decision process (the usual setting for reinforcement learning), all the information about the environment must be contained in the current state. This includes the time-step, if it affects any aspect of the environment.

For example, if an episode is terminated with probability 1% each step, then it's valid not to include the time-step as part of the state. On the other hand, if the episode terminates with probability n% after the nth step, then to have a correct MDP, the time-step should be part of the step.

In practice, for many environments such as Atari games or robotic manipulation tasks, people will cut off any episodes that go over some threshold number of steps, just to speed up training. They don't put the timestep into the state either, which is strictly speaking a violation of the markov property, but in practice it's not a huge deal and everything trains fine regardless.

In practice what the Q-network would learn in practice under this "stop the game above a certain timestep" strategy is probably something approaching

$$Q(s,a) = \sum_t P(t|s)Q'(s',a)$$

where $Q'$ is the actual state-action value function, and $Q$ is our function approximation. $s'$ is the actual state consisting of $(s,t)$. If $P(t|s)$ is a very "sharp" distribution (that is, the state gives a lot of information about which time-step it is), then there is no need to input the time-step separately.

On the other hand, if $s$ does not give a lot of information to $t$, then we might just have something approaching

$$Q(s,a) = \frac{1}{T}\sum_t Q'(s',a)$$

So the network would guess at the average expected reward over all time-steps, which shouldn't cause any big problems.

Another reason that this usually isn't a big deal is that for discount values such as 0.9, the reward decays pretty quickly, which means the $Q$-net doesn't care too much what happens more than 10 or 20 steps in the future. If we terminate all games over 1000 steps, then the base probability that the agent will be terminated in the next 10 or 20 steps is rather small, and shouldn't affect the predicted $Q$ much.