Long N-step Surrogate Stage Reward (LNSS)

Topics: Reinforcement Learning

Introduction

Long N-step Surrogate Stage (LNSS) reward is a suggested stage reward formulation aimed at reducing the variance of Q value in deep reinforcement training. It is a combination of single-step and n-step methods, in the sense that it uses a weighted average of the rewards of the next N steps in order to perform single-step bootstrapping. Because of that, it can be easily piggybacked on state-of-the-art algorithms like DDPG, TD3 and D4PG.

The limitations that LNSS addresses are:

• the high variance problem in training, which impedes successful training and reduces performance
• n-step methods with large $n$ result in a large reward scale with hinders training
• the performance of single-step and n-step methods degrades in environments with noisy reward signals or corrupted sparse rewards.

Mathematical Details

The concept of LNSS is to replace the single-step stage reward with a weighted average of the rewards of N steps. This way, the single-step bootstrapping method of DDPG and TD3 can be applied, but with a substitute stage reward computed from N steps.

Consider the N-step infinite horizon discounted return:

$$G_k = \sum_{t=k}^{k+N-1} \gamma^{t-k} \, r_t$$

Then, the LNSS substitute stage reward is as follows (1):

$$r'_{k} = \frac{\sum_{t=k}^{k+N-1} \gamma^{t-k} \, r_t}{ \sum_{n=0}^{N-1} \gamma^n}$$

Which is the weighted average of the rewards obtained over $N$ steps.

The equation can be further simplified by using the closed form of the geometric sum in the denominator:

$$r'_{k} = G_{k} \; \frac{\gamma - 1}{\gamma^N - 1}$$

The LNSS stage reward $r'_k$ for step $k$ can only be calculated at step $k+N-1$. Once $r'_k$ is obtained, the training tuple $(s_k, a_k, r'_k, s_{k+1})$ is added to the replay buffer $\mathbb{D}$. Training then proceeds as normal for the algorithms, for example by sampling a learning batch from the replay buffer and updating the neural network weights.

Because LNSS uses a weighted average and is used in single-step bootstrapping, $N$ can be increased to high values like $50$ or $100$ without affecting the reward scale. Also, in sparse environments, LNSS continuously and progressively provides a reward starting from N steps backward from the time of achieving the desired state. This reduces the sparsity of the environment and turns the sparse reward into a dense reward.

Pseudocode

References

1. Zhong, Junmin, Ruofan Wu, and Jennie Si. "A Long N-step Surrogate Stage Reward for Deep Reinforcement Learning." Advances in Neural Information Processing Systems 36 (2023): 12733-12745.