How would one normalize observations in off-policy online reinforcement learning?

Question

In off-policy algorithms such as DQN, you need to feed your input to a network twice. 1. When inputting into a network for predicting the Q values. 2. When feeding the input from the buffer to the network, to train it. If one doesn't have a standard way of normalizing, this would lead to a difference in the distribution between the inputs for prediction and training, despite having the same values.

My questions follows from the question that I asked over here (I have been told to use a new post for each question). The accepted answer suggested computing running statistics for my inputs. But I doubt that'll work with the problem described above.

Answer 1

(I'll provide a couple of ideas, but I don't think these would fully address the problem.)

If you normalize with running statistics (e.g. mean and std) as in the online case, you would get two different results for the same input, say $s$, because it would be normalized with different statistics. I think this would introduce some non-stationarity in the problem, which one usually wants to avoid.

One possible solution is to also store the statistics computed at timestep $t$, say $\mu_t$ and $\sigma_t$ such that when you replay a transition $(s_t, a_t, r_t, s'_t, \mu_t,\sigma_t)$ you also have the correct statistics for that timestep: which should avoid the distribution shift.

In general, if we assume that the distribution of $s_t$ is stationary i.e. it doesn't change with time, then, in the limit $t\to\infty$ (also assuming unbiased estimators), the statistics $\mu_t$ and $\sigma_t$ will converge to the true mean and std. If so, i.e. when you detect that $\mu_t$ is close to $\mu_{t+1}$ then you can stop storing the stats for each $t$ and only use the latest ones to avoid the extra memory cost. But storing that especially at the beginning of training can help stabilize the agent.

Alternatively, you can consider $K$ steps of interaction and keep the stats constant for that number of steps. Then you can either recompute them considering the new $K$ transitions or flush the replay buffer (which probably has the same effect, but reduces sample-efficiency.)