Time-Efficient Reinforcement Learning with Stochastic Stateful Policies

We very much appreciate the valuable suggestions of the reviewer and are especially grateful for the time invested by the reviewer in reading also the appendix. We agree with the reviewer that the old experiment section was not self-contained. In the updated version of the paper, we updated the experiment section to include 'at least a teaser' of all experiments; an idea from the reviewer we very much liked. Now. the reader gets an overview of all experiments without the need to dive into the appendix.

We have highlighted all changes in the paper in blue to make it easier for the reviewer to spot the differences. In the following, we answer to the specific concerns of the reviewer.

The Entire Theoretical part, especially the equations, should be formatted more concisely; there is no reason to display the equations in two lines. Notably, only only circa. 1.5 pages in the main part are left for the experimental results. Only two experiments are presented, and there are a few interesting ones in the appendix. The 'Stateful policies as inductive biases' experiment is especially interesting, and it is also mentioned in the introduction.

We updated the paper and shortened all two-line equations to single-line equations in the main part. We also shortened the theoretical part to make more space for experiments. We included a new figure, which gives an overview of all experiments. Every experiment is now discussed in a separate paragraph, including the inductive bias experiment. In doing so, all the highlights are in the main paper. If the reviewer thinks that there is still an important experiment missing in the main paper from the appendix, we would be happy to add it.

The main advantage of the approach over BPTT - improved computational efficiency is not shown in the paper. It would strengthen the paper when the computational benefits are demonstrated, using at least a table with actual wall-times comparison. It is only mentioned in the paper without proof: "The overall results show that our approach has major computational benefits w.r.t BPTT with long histories at the price of a slight drop in asymptotic performance for SAC and TD3".

The computation times were shown in Figure 2 (now Figure 3) together with the rewards. The diagram is labeled with 'Training Time [h]'. We agree that this was not clear from the main text in the old version. We now specifically refer to it in the main text. The training time plots show the time needed for 1 million steps averaged across all POMDP Gym tasks. We now also added detailed tables with specific training times for each environment in the Appendix. Note that the differences in training times in the imitation learning setting are comparable to the reinforcement learning results, as LS-IQ is based on SAC and Gail is based on PPO.

Often, results are mentioned in the main part with a reference to figures in the appendix. I emphasize that most readers will likely stop reading after the main part, so it would be nice to have at least a 'teaser' of the results in the main part.

All experiments from our paper are now presented in the main part. We absolutely agree with the reviewer that the experiments in the main sections were not self-contained before. We now give an overview of all tasks with a new figure and explain every experiment in a separate paragraph. In doing so, we reduced the amount of references to the appendix and made the main paper more self-contained.

Provide a wall-time comparison of the introduced S2PG approach with (as I understand more time-consuming) BPTT approach.

As mentioned before, this is now more clear in the paper.

Used assumptions in the theoretical part - argue their physical motivation and preferably cite some established works that introduced them

The assumptions used in the variance analysis are taken from prior work [1] and extended to the matrix setting. They basically assume a well-behaved policy structure by bounding the maximum magnitude of the gradients. The assumptions on the MPD used in the proofs of the policy gradient theorems are directly taken from prior work [2]. Similarly, they assume that the MPD is well-behaved, resulting in an MDP without singularities and/or explosion. We added the respective citations in the paper.

[1] Tingting Zhao, Hirotaka Hachiya, Gang Niu, and Masashi Sugiyama. Analysis and improvement of policy gradient estimation. In Proceeding of the Twenty-fifth Conference on Neural Information Processing Systems, Granada, Spain, December 2011.

[2] David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. De- terministic policy gradient algorithms. In Proceeding of the International Conference on Machine Learning, Beijing, China, June 2014.

We very much appreciate the detailed comments and suggestions of the reviewer. We have updated the main paper to address the concerns of the reviewer. We also added more detailed results, such as tables with explicit runtimes of all approaches in all environments in the appendix.

We have highlighted all changes in the paper in blue to make it easier for the reviewer to spot the differences. In the following, we answer the reviewer's specific questions.

Many details, such as the results for the memory task, are relegated to the appendix. Nonetheless, I do not regard this as a significant weakness, given that the main text is already very dense with foundational results.

We admit that the experiment section was not self-contained. We updated the paper now to give an overview of all experiments in a new Figure. Every experiment is now also getting a dedicated paragraph. Hence, the memory task itself and some of the results are now also shown in the main paper. In doing so, we reduced the amount of references to the appendix, making the experiment section easier to read.

Only ten seeds are utilized for the experiments, although it is well-known that MuJoCo tasks are prone to considerable variances in performance. I would recommend increasing the number of seeds to twenty for the final evaluation.

We agree with the reviewer that 10 seeds might not be enough. We are currently running all experiments again to increase the number of seeds to 25 for all experiments in the final version of the paper.

The paper would benefit from including an analysis of the variance of the gradient for both the proposed method and BPTT, even if on a very simple benchmark. Additionally, it would be beneficial to examine the issues of vanishing and exploding gradients for both BPTT and the proposed method in at least one benchmark.

This is indeed a very good idea. We are currently working on a toy example to empirically evaluate the variance of both gradient estimators in the policy gradient setting. We believe that this is a great addition to the theoretical analysis done in the paper.

The occupancy measure is introduced without defining z

Thanks for pointing this out! This is fixed in the updated paper.

Is the initial internal state learned, or is it initialized to zero at the beginning of each episode?

We always initialize the policy state to zero. We made this more clear in the updated paper. However, the initial state of the policy is a great way to induce any kind of knowledge a priori into the policy. For RNNs, we are not aware of a useful initial state, but in the case of an ODE, there are many interesting initial states you might want to try out.

Does the limit of the stochastic policy gradient converge to the deterministic policy gradient when the variance of the action and subsequent internal state approaches zero? In other words, is there a result analogous to Theorem 2 in the Deterministic Policy Gradient paper?

This is a very good question. While we genuinely believe that a similar Theorem exists for our stateful policy gradient, this requires a proof. As the proof of Theorem 2 in the Deterministic Policy Gradient paper is non-trivial, we need time to carefully do the proof. We will try to do this for the final version of the paper.

We thank the reviewer again for his suggestions and very much appreciate the time and effort spent.

We very much appreciate the detailed comments of the reviewer. We have updated the main paper to address the reviewer's concerns.

Note that we additionally improved the experiments section by making the latter more self-contained and less reliant on the appendix. We also added more detailed results, such as tables with explicit runtimes of all approaches in all environments, in the appendix. We have highlighted all changes in the paper in blue to make it easier for the reviewer to spot the differences. In the following, we answer to the specific concerns of the reviewer.

I am not convinced by the arguments of the authors. In their formulation, even if using stochastic states prevents one from computing an analytic gradient, I don't understand why it would solve the problem of capturing long-term information. Indeed, when computing $p(a_t,z_t|s_t,z_{t-1})$, then this probability depends on the previous timestep, and so on, such that finding a good solution to the problem would need to propagate the loss to the previous timesteps to capture a good sequence of states. So there is still backpropagation through time, even if it is not made by the analytical gradient.

In the BPTT algorithm, the gradients of the policy are computed analytically for the whole history up until time $t$. In contrast, our approach uses stochastic exploration to estimate the stateful policy gradient locally at time $t$ given an estimate of the cumulative return. This is basically the same principle as applying the likelihood-ratio trick used in standard reinforcement learning to estimate the gradient of a policy in an unknown environment. In reinforcement learning, it is not possible to estimate the gradient of the policy without sufficient exploration. We follow the same principle to estimate the gradient of a stateful policy without relying on the analytical gradient of the history.

While it is possible to interpret our approach as backpropagating information implicitly using our stochastic gradient estimator, we prefer the forward-looking perspective (the expected cumulative reward) used in reinforcement learning. This forward perspective becomes even more prominent when looking at our action value $V(s, z)$, i.e., the expected cumulative when being in the states $s$ and $z$, and sampling $a$ and $z'$ from our policy from there on. In other words, our policy needs to learn in the early states of a trajectory what information to encode for the future rather than propagating information back in time.

It is also noteworthy that we show that the gradient of the policy can also be computed deterministically if we have access to a critic $Q(s, a, z, z')$. However, to estimate this critic, exploration is still necessary. With the conventional critic $Q(s, a)$, BPTT would still be necessary to propagate the information back in time.

Then, usually, relying on stochastic variables decreases the sample efficiency of training methods. This is why people are using for example the reparametrization trick that allows one to compute an analytical gradient over a stochastic variable, to speed up training. Here the authors are claiming the opposite. So there is one point that I didn't catch in this paper, and I would like the authors to better explain why using stochastic nodes would avoid the problem of propagating information to the previous timesteps, and why they would expect a better sample efficiency than using an analytical gradient

We would like to point out that our approach is not more sample-efficient than BPTT. In fact, we agree with your statement that the efficiency in samples is generally lower compared to backpropagation through time, and this is something we see in the empirical results (robust RL results). What we claim is that our gradient estimator is more time-efficient than BPTT as the calculation of the gradient of a trajectory is an inherently sequential computation. This is particularly evident when taking the full history to compute the gradient. While BPTT can become more time-efficient by truncating the history, this leads to a biased gradient estimate. In contrast, our approach always provides an unbiased gradient estimate. Also, our gradient estimator is less dependent on the architecture of the policy, making it applicable to any stateful policy (c.f., theoretical analysis). For BPTT, specialized architectures like GRUs and LSTMs were introduced to cope with exploding and vanishing gradients. We see better efficiency in samples only when scaling the complexity of the tasks (Ant, and Humanoid in the robust RL, and all imitation learning experiments). We trace this back to the high dimensionality of the state space, which makes BPTT harder to optimize.