How to deal with episode termination in Advantage Actor-Critic algorithm?

Question

Advantage Actor-Critic algorithm may use the following expression to get 1-step estimate of the advantage:

$ A(s_t,a_t) = r(s_t, a_t) + \gamma V(s_{t+1}) (1 - done_{t+1}) - V(s_t) $

where $done_{t+1}=1$ if $s_{t+1}$ is a terminal state (end of the episode) and $0$ otherwise.

Suppose our learning environment has a goal, collecting the goal gives reward $r=1$ and terminates the episode. Agent also receives $r=-0.1$ for every step, encouraging it to collect the goal faster. We're learning with $\gamma=0.99$ and we terminate the episode after $T$ timesteps if the goal wasn't collected.

For the state before collecting a goal we have the following advantage, which seems very reasonable: $A(s_t,a_t) = 1 - V(s_t)$.

For the timestep $T-1$, regardless of the state, we have: $ A(s_{T-1},a_{T-1}) = r(s_{T-1}, a_{T-1}) - V(s_{T-1}) \approx -0.1 -\frac{-0.1}{1-\gamma} = -0.1 + 10 = 9.9 $ (this is true under the assumption that we're not yet able to collect the goal reliably often, therefore the value function converges to something close to $\frac{r_{avg}}{1-\gamma} \approx -10 $ ).

Usually, $T$ is not a part of the state, so the value function has no way to anticipate the sudden change in reward-to-go. So, all of a sudden, we got a (relatively) big advantage for the arbitrary action that we took at the timestep $T-1$. Following the policy gradient rule, we will significantly increase the probability of an arbitrary action that we took at the end of the episode, even if we didn't achieve anything. This can quickly destroy the learning process.

How do people deal with this problem in practice? My ideas:

1. Differentiate between actual episode terminations and ones caused by the time limit, e.g. for them we will not replace next step value estimate with $0$.
2. Somehow add $t$ to the state such that the value function can learn to anticipate the termination of the episode.

As I noticed, the A2C implementation in OpenAI baselines does not seem to bother with any of that:

• A2C implementation - calculation of discounted rewards

• discount_with_dones function

Answer 1

First a small note: I don't think your expression for $A(s_{T-1}, a_{T - 1})$ looks correct. If we assume that $V(s_T) = 0$ (i.e., assume that we cannot possibly reach the goal in one single step from $s_{T - 1}$), we have:

\begin{align} A(s_{T-1}, a_{T-1}) &= r(s_{T-1}, a_{T-1}) + \gamma V(s_{T}) (1 - done_{T}) - V(s_{T-1}) \\ &= r(s_{T-1}, a_{T-1}) - V(s_{T-1}). \end{align}

In this expression, we'd normally have that $r(s_{T-1}, a_{T-1}) = -0.1$, whereas you seem to have mistakenly taken $+0.1$ in your post.

---

Those details aside; yes, there can be sudden bumps of positive advantage estimates as you described. This is not a problem though, this is exactly what we'd expect to happen given your description of the environment.

You describe an environment in which the agent is likely going to be wandering around randomly (at least when it hasn't learned a good policy yet), and incurring negative rewards over and over again. This naturally leads to negative value estimates for all encountered states. Suddenly, it does something and the episode terminates; it receives a nice reward of $0$ rather than yet another negative reward (this actually "feels" like a bonus, a positive reward, something more than was expected). When your agent has not yet learned a good policy that can reach the better reward of $1$, this is indeed a good result, a good action, and it rightfully should get reinforced.

Because this event of the episode terminating is mostly uncorrelated with the state (I say "mostly", because in theory it probably ends up being a slightly rarer event in states close to the goal than in states far away from the goal), it will eventually (after sufficient training time) end up occurring approximately equally often in all states. From the perspective of an agent that is oblivious to the current time step, this will be perceived as an event that can simply occur by pure chance in a nondeterministic environment.

This is not necessarily a problem. It can slow down learning due to increased variance in your reward observations (which can be addressed by using low learning rates / large batch sizes), but Reinforcement Learning algorithms are almost always naturally built to handle nondeterministic environments, it can work this out, it can average out all the different outcomes observed for the same state+action pairs. This is not a problem that requires dealing with.

---

My ideas:

1. Differentiate between actual episode terminations and ones caused by the time limit, e.g. for them we will not replace next step value estimate with $0$.
2. Somehow add $t$ to the state such that the value function can learn to anticipate the termination of the episode.

The first idea fundamentally changes the quantity that your algorithm is learning, it will essentially make your learning algorithm incorrect. There's always a chance that it might still appear to learn something useful (many Machine Learning/Reinforcement Learning algorithms can still appear to be okay even when there are bugs/technically incorrect parts), but it'll very likely perform worse.

The second idea, while not necessary as explained in my answer above, may still be beneficial to learning speed provided it is done well. It may help because it can add to the power of your algorithm to "explain" its observations, and more importantly "explain" the variance in its observations.

The main problem I see with adding $t$ to your input is that it is not naturally a binary variable. Very often you'll find that we're just using a bunch of binary inputs in (Deep) Reinforcement Learning algorithms. When all inputs are of the same magnitude like that, it tends to be easier to get the learning process to run well, tune hyperparameters like learning rate, etc. If you suddenly plug in an additional input which can take significantly larger values (like straight up adding $t$ as an input), this will be more difficult. Adding $\frac{t}{T}$ as an input may be better, since that will always still be bounded between $0$ and $1$.

Answer 2

I would like to point out this paper: https://arxiv.org/pdf/1712.00378.pdf, that answers exactly that question.

We then showed that, when learning policies for time-unlimited tasks, it is necessary for correct value estimation, to continue bootstrapping at the end of the partial episodes when termination is due to time limits, or any early termination causes other than the environmental ones.

The authors therefore argue, that the target $y$ for a one-step TD update, after transitioning to a state $s′$ and receiving a reward $r$, should be:

$$ y = \begin{cases} r & \text{if environment terminates}\\ r+\gamma\hat{v}_{\pi}(s') & \text{otherwise (including timeouts)} \end{cases} $$

Answer 3

The practical need for a $T$ or $t_{max}$ is based on the importance of prudent doubt as a cost containment measure. At a certain point, it is wise to abandon an objective the resolution of which cannot be certain. This is a technique used in the minimization of a much wider loss function associated with the business into which the seeking of the learning objective fits.

Dismissing the easily retained knowledge of the state of the independent process variable $t$ in the context of its domain $[1, T]$ as irrelevant is like working on a project without project management. What things to skip and what risks to take cannot be accurately evaluated, reducing the probability of success in relation to its optimum.

Certainly the value gained at any step is an indicator of the advantage of the current path, an affirmation that impacts the probability of achieving favorable results before $t$ reaches $T$. Certainly any knowledge that can be used to predict possible values along the paths that begin with choices that can be made at any step should impact that choice.

The question author's intuition that the critic should be aware that the stakeholder may pull the plug on the project soon is correct. In terms of the mathematics, whether urgency can be formulated as something like $u(t) \propto \dfrac {1} {T - t + 1}$ and used as a term or a factor requires some thought about the probabilities involved. The advantage expression should reflect the application of the basic relations of probability to the probability distributions as seen considering all the information known at any temporal point along the pursuit of the objective.

The only occasion when a variable that can be known should be excluded from the advantage expression is if it can be shown to be absolutely irrelevant or that its relevance is sufficiently small to warrant exclusion to gain computational speed in exchange for loss of accuracy and/or reliability.