Time-Efficient Reinforcement Learning with Stochastic Stateful Policies
Authors: Firas Al-Hafez
Guoping Zhao
Jan Peters
Davide Tateo
Venue: ICLR 2024 poster (2024)
Keywords: reinforcement learning
recurrent neural networks
stateful policies
backpropagation through time
imitation learning
Primary area: reinforcement learning
TL;DR: Our novel gradient estimator makes training of stateful policies simpler and faster.

Abstract: Stateful policies play an important role in reinforcement learning, such as handling partially observable environments, enhancing robustness, or imposing an inductive bias directly into the policy structure. The conventional method for training stateful policies is Backpropagation Through Time (BPTT), which comes with significant drawbacks, such as slow training due to sequential gradient propagation and the occurrence of vanishing or exploding gradients. The gradient is often truncated to address these issues, resulting in a biased policy update. We present a novel approach for training stateful policies by decomposing the latter into a stochastic internal state kernel and a stateless policy, jointly optimized by following the stateful policy gradient. We introduce different versions of the stateful policy gradient theorem, enabling us to easily instantiate stateful variants of popular reinforcement learning and imitation learning algorithms. Furthermore, we provide a theoretical analysis of our new gradient estimator and compare it with BPTT. We evaluate our approach on complex continuous control tasks, e.g. humanoid locomotion, and demonstrate that our gradient estimator scales effectively with task complexity while offering a faster and simpler alternative to BPTT.

Published as a conference paper at ICLR 2024
TIME-EFFICIENT REINEORCEMENT LEARNING WITH STOCHASTIC STATEFUL POLICIES
Firas Al-Hafez', Guoping Zhao Jan Peters Davide Tateo
1 Intelligent Autonomous Systems, 2 Locomotion Laboratory
German Research Center for AI (DFKI), Centre Cognitive Science, Hessian.Al
TU Darmstadt, Germany
{name. wurnameeetu-deeee
de
ABSTRACT
Stateful policies play an important role in reinforcement learning, such as handling partially observable environments, enhancing robustmess, or imposing an induc- tive bias directly into the policy structure. The conventional method for training stateful policies is Backpropagation Through Time (BPTT), which comes with significant drawbacks, such as slow training due to sequential gradient propaga- tion and the occurrence of vanishing or exploding gradients. The gradient often truncated to address these issues, resulting in a biased policy update. We present a novel approach for training stateful policies by decomposing the latter into a stochastic internal state kernel and a stateless policy, jointly optimized by follow- ing the stateful policy gradient. We introduce different versions of the stateful policy gradient theorem, enabling us easily instantiate stateful variants of pop- ular reinforcement learning and imitation learning algorithms. Furthermore, we provide theoretical analysis of our new gradient estimator and compare it with BPTT. We evaluate our approach on complex continuous control tasks, e.g. hu- manoid locomotion, and demonstrate that our gradient estimator scales effectively with task complexity while offering faster and simpler alternative to BPTT
INTRODUCTION
Stateful policies are a fundamental tool for solving complex Reinforcement Learning (RL) problems. These policies are particularly relevant for RI. in a Partially Observable Markov Decision Process (POMDP), where the history of interactions needs to processed at each time step. Stateful poli- cies, such as Recurrent Neural Network (RNN), compress the history into a latent recurrent rep- resentation, allowing them to deal with the ambiguity environment observations, Ni et al. (2022) have shown that RNN policies can be competitive and even outperform specialized algorithms in many POMDP tasks. Besides POMDPs, stateful policies can be used to incorporate inductive biases directly into the policy, e.g. a stateful oscillator learn locomotion Ijspeert et al. (2007); Bellegarda & Ijspeert 2022a), or to solve Meta-RL tasks by observing the history of rewards Ni et al. (2022).
Existing methods for stateful policy learning mostly rely either black-box/evoluttonaryy optimiz- ers or on the BPTT algorithm. While black-box approaches struggle with high-dimensional pa- rameter spaces, making them less suitable for neural approximators, BPTT has shown considerable success in this domain [akker (2001): Wierstra et al. 10); Meng et al. (202I). However, BPTT suffers from significant drawbacks, including the need for sequential gradient propagation through trajectories, which significantly slows down training time and can lead to vanishing or exploding gradients. To address these issues, practical implementations often use a truncated history, intro- ducing a bias into the policy update and limiting memory to the truncation length. Furthermore, integrating BPTT into standard RI. algorithms is not straightforward due to the requirement of han- dling sequences of varying lengths instead of single states.
In this paper, we propose an alternative solution to compute the gradient a stateful policy. We decompose the stateful policy into stochastic policy state kernel and conventional stateless pol- icy, which are jointly optimized by following the Stochastic Stateful Policy Gradient (S2PG). By
Thee code is available httpp:://gttbbb.com//

---

Published as a conference paper at ICLR 2024
adopting this approach, we can train arbitrary policies with internal state without BPTT. S2PG not only provides an unbiased gradient update but also accelerates the training process and avoids issues like vamishing exploding gradients. In addition, our approach can be applied to any exist- ing RL. algorithm by modifying a few lines of code. We also show how to extend the S2PG theory to all modern RL approaches by introducing different versions of the policy gradient theorem for stateful policies. This facilitates the instantiation of stateful variations of popular RL algorithms like Soft-Actor Critic (SAC) Haarnoja et al. (2018), Twin Delayed DDPG (TD3) Fujimoto et al. (0118), Proximal Policy Optimization (PPO) Schulman et al. 20177) as well as Imitation Learning (IL) algorithmss like Generative Adversarial Imitation Learning (GAIL) HHo & Ermon (2016) and Least-Squares Inverse Q-Learning (LS-IQ) Al-Hafez et al. (2023). Our approach is also applicable to settings where the critic has only access to observations.
such cases, we combine our method with Monte-Carlo rollouts, as typically done in PPO. In the off-policy scenario, we use a critic with privileged information Pinto et al. (2018): Lee et al. (2020b); et al. (2018).
To evaluate our approach, we conduct theoretical analysis on the variance of S2PG and BPTT, introducing two new bounds and highlighting the behavior of each estimator in different regimes. Empirically, compare the performances in common continuous control tasks within POMDPs when using RL. To illustrate the potential use of S2PG for inductive biases in policies, we give an example of trainable Ordinary Differential Equation (ODE) of an oscillator used in the policy to- gether with an Feed-Forward Network (FFN) to achieve walking policies under complete blindness. Finally, We also demonstrate the scalability of our method by introducing two complex locomo- tion tasks under partial observability and dynamics randomization. Our results indicate that these challenging tasks can be effectively solved using stateful IL algorithms, and demonstrate that S2PG offers simple and efficient alternative to BPTT that scales well with task complexity.
Related Work. RNNs are the most popular type of stateful policy. Among RNNs, gated recur- rent networks, such as Long-Short Term Memory (LSTM) [Hochreiter & Schmidhuber (1997) or Gated Recurrent Units (GRU) Cho et al.
(014) networks, are the most prominent ones Ni et al.
(22T) Heess et al. (2015); [Espehol et al. (2018): Yang & Nguyen (2021). Wierstra et al. (2010) introduced the recurrent policy gradient using LSTMs
(2022); Wierstra et al. (2010)g Meng et al.
and the GPOMDP Baxter & Bartlett (2001) algorithm, where BPTT was used to train the trajectory. Many subsequent works have adapted BPTT-based training of RNNs to different RL algorithms. Heess et al. (015)) introduced the actor-critic Recurrent Deterministic Policy Gradient (RDPG) and Recurrent Stochastic Value Gradient (RSVG), which were later updated to the current state- of-the-art algorithms like Recurrent Twin-Delayed Deep Deterministic Policy Gradient (RTD3) and Recurrent Soft Actor-Critic (RSAC) Yang & Nguyen (02TT. Many works utilize separate RNNs for the actor and the critic in RTD3 and RSAC, as yields significant performance gains compared to shared RNN Ni et al. (202);; ; Meng et al.] (2021); Heess et al. (2015): Yang & Nguyen (2021). Although LSTMs have a more complex structure, GRUs have demonstrated slightly better perfor-
mance in continuous control tasks Ni et al. (022). Stateful Policies can also be used to encode an inductive bias into the policy. For instance, Central Pattern Generator (CPG) is a popular choice for locomotion tasks [jspeert] (2008); [Bellegarda & Ijspeent (2022b); Campanaro et al. (202T). I). We conduct experiments with CPGs in Section [..] where we also present relevant related work. Us- ing similar approach to the one presented in this paper, Zhang et al. 2016) introduce a stochastic internal state transition kernel. In this approach, states are considered as a memory that the policy can read and write. However, they did not extend their approach to the actor-critic case, which is the main focus of our work. |Rakelly et al. (2019) proposed a Meta-RL. approach to learn a distribution over latent variables that encodes history to extract context variables. Similarily, other approaches such as DVRL [gl et al. (2018) and SLAC Lee et al. (2020a) force the recurrent state to be belief state, exploiting learned environment model. An alternative to stateful policies, particularly useful in POMDPs settings, is to use history of observations. These policies use either fully connected MLPs, time convolution Lee et al. (2020b), or transformers Lee et al.] (2023); [Radosavovic et al. 2023). It can be shown that for some class of Markov Decision Processs (MDPs), these archi- tectures allow learning near-optimal policies |Efroni et al. (2022). The main drawback is that they require high dimensional input, store many transitions in the buffer, and cannot encode inductive biases in the latent space. Finally, the setting where privileged information from the simulation is used is widespread in sim-to-real robot learning Lee et al.) (2020b); |eng et al. (2018). Compared to so-called teacher-student approaches Lee et al. (2020b), which learn a privileged policy and then train a recurrent policy using behavioral cloning, our approach can learn a recurrent policy online using privileged information for the critic.
2

---

Published as a conference paper at ICLR 2024
2 STOCHASTIC STATEFUL POLICY GRADIENTS
(rr,y,)) 7,4, where S is the state space, A is the action space, P: 8 x A S - R + is the transition kernel, r x A - R is the reward function, y is the discount factor, and L S - R is the initial state distribution. At each step, the agent observes a state s E S from the environment, predicts an action transitions with probability P(s'ls, a) into the next state s' E S. where it receives the reward r(ss, a). We define an occupancy measure pa ( z) Er= 7 Pr(s 8t A 2 giving metric on how frequently the tuple (s, is visited, where 2 is the internal state of the policy. The POMDP tuple (S,A, O, P.O. r. 7, s) extending the MDP by the observation space O and the conditional probability density O O S + IR+, such that for every state 8 E S, the probability density of o E O (olss) trajectory is a (possibly infinite) sequence of states, actions, and observations sampled by interacting with the environment under a policy æ. The return of a trajectory is defined as (T) = Ei=o 7'r(st, a.). We define the observed history h4 = (su....82. ) as the sequence of states taken up to the current timestep t. In the partially observable case, the history is given by observations and actions such that ht (oo, an 00-1, a,-1, 0z). In general, an optimal policy for POMDP depends the full observed history 14, i.e. x (-11,). However, most of the time, we can substitute h, with sufficient statistics 2; = S(he), namely the internal state of the policy.
Preliminaries. An MDP is tuple
A using the policy π S x A Rt, and
2.1 POLICY GRADIENT OF STATEFUL POLICIES
Wierstra et al. (2010) presented the first implementation of the policy gradient for recurrent policies. This formulation based on the likelihood ratio trick, i.e. the score function estimator. To obtain this formulation, the authors consider generic history-dependant policy vaali,,, obtaining the following gradient formulation
Vo = Vo log (a (o((
(1)
Note that each Ve log vo Poraell (4.) term in equation I depends on the whole history h, at each timestep t which makes BPTT necessary for stateful policies, such as RNNs. However, BPTT comes with many issues. Firstly, this method is inherently sequential, limiting the utilization of potential gradi- ent computations that could be performed in parallel. Secondly, when dealing with long trajectories, the gradients may suffer from the problems of exploding or vanishing gradients, which can hinder the learning process. Therefore, histories are often truncated, leading to a biased gradient estimate.
To solve these issues, consider different policy structure. Instead of looking at policies with an internal state, we model our policy as a joint probability distribution over actions and next internal states, i.e. (a, 2') To(-1s. (-|s 2). By considering stochastic policy state transition, we derive a policy gradient formulation that does not require the propagation of the policy gradient through time but only depends on the local information available, i.e. the (extended) state transition.
Lemma 2.1. (Stochastic Stateful Policy Gradient) Let ae(a, z/|ss, 2) be a parametric policy repre- senting the joint probability density used to generate the action and the next internal state z', and let be the extended trajectory including the internal states. Then, the S2PG can be written as
VoJ(mo) =
Vo log no (ar, 2441 188, 2x)
This lemma is straightforward to derive using the classical derivation of the policy gradient. The complete proof can found in Appendix in equation 2 depends the information of the current timestep I, i.e. dr, 1,41 and z1. We can deduce the following from Lemma 2.1
contrast to the gradient in equation the gradient
The gradient of stateful policy can be computed stochastically, by treating the internal state of policy as a random variable. This modification results in an algorithm that does not require BPTT, but only access to samples from
We allow to trade off speed and accuracy while keeping the gradient estimation unbiased. While truncating the history to save computation time biases the gradient, estimating the expectation with a finite amount of samples of z does not, In contrast to BPTT, this formulation supports par- allel computation of the gradient, allowing fully exploit the parallelization capabilities of modern
3

---

Published as a conference paper at ICLR 2024
2Q
80
Figure I: Stochastic computational graph illustrating, left to right, a comparison of stateless policy, stateful policy trained with BPTT, and stateful policy incorporating a stochastie internal state transition kernel. Deter- ministic nodes (squares) allow for the passage of analytic gradients, while stochastic nodes (circles) interrupt the deterministic paths. The input node to the graph is represented by 9. The blue lincs indicate the determin- istic paths for which analytic gradients are available for specific action. For detailed demonstration of how stochastic node gradients are calculated compared to deterministic ones, refer to Figure Appendix
automatic differentiation tools and parallel simulation environments, As the batch size for gradient computation increases, the variance of our gradient estimator decreases, as shown later Section
We present an alternative perspective by explaining the core principle of Lemma 2. in the context of stochastic computation graphs, as proposed by |Schulman et al. al.] (01T5)). The key idea of Lemma 2.1 is to shield the deterministic path through internal states using a stochastic node, specifically our stochastic internal state kernel. Figure [illustrates this by comparing analytic gradient paths of the internal state across different policies. While S2PG resolves issues with exploding and vanishing gradients in long sequences, introduces potential high variance due to the stochastic internal state kernel. To address this, propose compensating for the added variance by utilizing state-offin state-of-the-art tht-hhe-rt actor-critic methods while maintaining the same policy structure.
2.2 ACTOR-CRITIC METHODS FOR STATEFUL POLICIES
To define actor-critic methods, need to extend the definitions of value functions into the stateful policy setting. Let s be some state of the environment, z some state of the policy, a given action, and z the next policy state. We define the state-action value function of our policy π (a, 2 2/|8, z) as
("(,,,,a,=') = r(s,u) + +
E
[Q" = =s),,)) + (V"''''
(3)
with s' ~ P(.(s,a) and x") ~ π (-15/ 2/). Furthermore, we define the state-value function as [e 2s, where (a, ~ T-11, 2).
(=((, E, a."
Note that these definitions are equivalent to those of the value and action-valus functions in a Markov policy, Indeed, the value function ((s, =) the expected discounted return achieved by the policy starting from the state with initial internal state z. Thus, we can write the expected discounted return of stateful policy as [(x)= f (ss Vs, 2)dsdz with the initial state distribution ((s, 2).
The fundamental theorem of actor-critic algorithms is the Policy Gradient Theorem (PGT) Sutton et al. 1999). This theorem lays the connection between value functions and policy gradients. Many successful practical approaches are based on this idea and constitute fundamental part of modern RL. Using our stochastic internal state kernel and the definitions of value and action-value function in equation
4 and 3 we provide straightforward derivation of the PGT for stateful policies.
Theorem 2.2. (Stateful Policy Gradient Theorem) Lei 8Ta (s, 2) be the occupancy measure and
QTe (s, 2') be the value function of a parametric policy To Te(a, 2/|s,,)) then
VoJo = pT -
Sal Vonola, 2'18, 2,0, dz' da dz ds,
is the gradient of the stateful policy.
The proof follows the lines of the original PGT and is given in Appendix A.3 Differently from Sutton et al. (1999). which considers only the discrete action setting, we provide the full derivation for continuous state and action spaces.
4

---

Published as a conference paper at ICLR 2024
Theorem 2.2 elucidates the key idea of this work:
We learn a Q-function to capture the values of internal state transitions, compensating the additional variance in the policy gradient estimate induced by stochastic internal state transition kernel. This approach enables the efficient training of stateful policies without the need for BPTT, effectively reducing the computational time required train stateful policies to that of stateless policies. Figure i in Appendix
B further illustrates and reinforces this concept.
While the PGT is a fundamental building block of actor-critic methods, some of the more successful approaches, namely the Trust Region Policy Optimization (TRPO) [Schulman et al. (2015b) and the PPO Schulman al. (2017) algorithm, are based on the Performance Difference Lemma Kakade & Langford (2002). We can leverage the simplicity of our policy structure to derive Lemma A.1
modified version of the original Lemma for stateful policies, allowing us to implement recurrent versions of PPO and TRPO with S2PG. The Lemma and the proof can be found in Appendix A.4
2.3 THE DETERMINISTIC LIMIT
The last important piece of policy gradient theory is the definition of the Deterministic Policy Gra- dient Theorem (DPGT) Silver al. (014). This theorem can be seen as the deterministic limit of the PGT for some class of probability distributions. DPGT provides the basis for many successful actor-critic algorithms, such as Deep Deterministic Policy Gradient (DDPG) Lillicrap et al. (2016) and TD3 Fujimoto et al. (2018). We can derive the DPGT for our stateful policies:
[DDlG)Lill
Theorem 2.3. (Stateful Deterministic Policy Gradient Theorem) Let pa(8,2) = [,e, 1 oo be a deterministic policy, where v=.... = pā(s,2) 2) represent the action model and v2,2,0 uĀ (s, represents the internal state transition model. Then, under mild regularity assumptions, the policy gradient of the stateful deterministic policy under Ph: pp (s, 2) is
NeJ (Ha()))= Pt. Ve (1.20(()
[ 5,2
+aa. 2 ( (s,Z,)" )asto dsdz.
The proof is given in Appendix A.5 It is important to notice that, while the deterministic policy gradient theorem allows computing the gradient of a deterministic policy, assumes the knowledge of the Q-function for the deterministic policy 14g However, to compute the Q-function of given policy is necessary to perform exploration. Indeed, accurately estimate the Q-values, actions different from the deterministic policy needs to be taken. We want to stress that in this scenario, in- stead of using the standard Q-function formulation, we use our definition of Q-functions for stateful policies: therefore, must also explore the internal state transition. This key concept explains why in our formulation it is not necessary to perform BPTT even in the deterministic policy scenario.
2.4 THE PARTIALLY OBSERVABLE SETTING
Within this section, we introduce the POMDP settings used within this work and adapt the previous theory accordingly. Therefore, we extend Lemma 2.1 as follows.
Corollary 2.4. The stateful policy gradient in partially observable environments
VoI(To)
=
E oa al(e(
Σ Vo log ne(ae, 24 26 [(F)
The proof is given in Appendix A.2 Corollary 2.4 offers a particularly intriguing perspective: If we have knowledge of J(7), it possible to learn an internal state representation ž, that effectively encapsulates past information without looking back time, In other words, it enables to learn the compression of past information without explicitly considering previous time steps. This stands in stark contrast to BPTT, where all past information is used to learn a belief of the current state. Similarly Corollary 2.4 we can extend Theorem 2.2
( ( (00
/ VoTola, x/ 10, 2) (o, (s, a,2²) dads
VeJo
Js
is important to note the distinction between the Q-function, which receives states, and the pol- icy, which relies observations only. While it is possible to lear Q-function observations
5

---

Published as a conference paper at ICLR 2024
via bootstrapping, the task of temporal credit assignment in POMDPs is very challenging, often requiring specialized approaches Igl et al. (2018): [ee et al. (2020a). This difficulty stems from the following causality dilemma: On one hand, the policy can learn condensed representation of the history assuming accurate Q-value estimates. On the other hand, the Q-function can be learned us- ing bootstrapping techniques assuming reliable condensed representation of the history. Because both, bootstrapping a Q-value function and learning a representation of the history, rely on mean- ingful representations of the other, doing both at the same time often results in unstable training. We show in Section 3 that it is possible to combine Monte-Carke estimates with critic that uses observations only using PPO to learn in true POMDP setting. However, we limit ourselves to the setting with privileged information in the critic, as shown in Equation 5 for approaches that learn value function via bootstrapping.
VARIANCE ANALYSIS OF THE GRADIENT ESTIMATORS
To theoretically assess our estimator's variance, we provide the upper bounds of BPTT and S2PG in the Gaussian policy setting with constant covariance matrix E. As done in previous work Pork(apinii al. (202); Zhao al. (201T). we define the variance of a random vector as the trace of its covariance matrix (see equation equation 17 in Appendix (A.6). Furthermore, for both gradient estimators, we assume that the function approximator for the mean has the following structure
pa(st, 2x) = fo(st, (8 21) (8t, 21) (for S2PG) 21+1 (8, =) (for BPTT).
(6)
Hence, [0( 2+) constitutes the mean of the Gaussian for both S2PG and BPTT, while no ()1, z4) constitutes the internal transition function for BPTT and the mean of the Gaussian distribution of the S2PG.Figuree
hidden state in S2PG. Figure AppendixiDjiluustates the policy structure. Additionally, I consti-
tutes the covariance matrix for the Gaussian distribution of the hidden state in S2PG. Our theoretical analysis is based on prior work Papini et al. (2022); Zhao et al. (2011). which analyzes the variance of a REINFORCE-styL policy gradient estimator for stateless policies and a univariate Gaussian distribution. We extend the latter to policy gradient estimators using BPTT and S2PG, and broaden the analysis the more general multivariate Gaussian case allow investigation on arbitrarily large action spaces. In the following, we derive upper bounds for the variance of the policy gradient for BPTT and S2PG and highlight the behavior of each estimator in different regimes. Therefore, we exploit the concept of the Frobenius norm of a matrix ||~A||; and consider the following assumptions
Assumption 3.1. r(s, as),,,8') 8') [-R, R] for R > O and Vs,, 2;
Assumption 3.2. zy Jn(s, z:) 11, - F, ||
2 5e (M),, I), & H, 11 Jo x)I,, S K Vss, 22
Assumption 3.3. || o2. (s: 24 11 Z
Vs..
extended from Zhao et al.
2011). Then, the upper bound on the variance of BPTT is given by:
Theorem 3.4. (Variance Upper Bound BPTT) Let vo h4) be a Gaussian palicy of the form N N( ((e)), Σ) and let T be the trajectory length. Then, under the assumptions 3.1 3.2 and 3.3
the upper bound on the variance of the REINFORCE-stye policy gradient estimate using BPT IS
R2 |=212== 1|
(1 1 -2
Var [Vo
5
(TF² 2FHKŽ + =H
N(1 - 7)²
where we define 2 = Ei-lz-i-> and Z = Σt- ( z++r)) for brevity.
Similarly, the upper bound on the variance of S2PG is given by:
Theorem 3.5. (Variance Upper Bound S2PG) Let To(ar, 244//st, 24) be Gaussian policy of the form N(ar. 2441|po 21), E, Y) and let T be the trajectory length. Then, under the assumptions
3.1 and] 3.2 the upper bound on the variance of the REINFORCE-Styl policy gradient estimate
using S2PG is
R² R 11
72 7 (TF" +TH
Var [oJsrcco)) s
(8)
N(1 7)²
Asgru
The proof of both Theorems can be found in Appendix A.6.1 and A.6.2 When looking at Theorem 3.4 and 3.5 it can be seen that both bounds consist of constant factor in front and sum over

---

Published as a conference paper at ICLR 2024
Robust RL Tasks
Memory Tasks
Task
Randomiy selected
Parametric
At
2t-1
Bited Area
FFN
Observbbe Area
handemty celected inittal Postions
Imitation Learning Tasks
From left to right, Gym locomotion tasks with randomized masses and hidden velacities, memory tasks with information hidden outside the observable area, and an ODE tasks where the policy is blind and the parameters of the ODE and FFN are Icarned. Bottom: POMDP Humanoid walking/running task at 4.5 km/h respectively 9 km/h. The different humanoids should resemble from left right an adult, a teenager (~112 years), a child (~55 years), and a toddler (~1-2 years). The target speed scaled for the smaller humanoid according to size. The type of humanoid hidden to the policy. single policy is able to Icarn all gaits. On the right, random carry-weight task using the Atlas humanoid. The weight randomly sampled between 0.1 and 10 kg at the beginning of the episode and is hidden the policy.
Figure 2: Overview tasks. Top:
F², which is a constant defining the upper bound on the gradient of (s;. (.). Additionally, both bounds introduce an additional term, which consider as A. In fact, the latter is a factor defining
the additional variance added when compared to the variance of the stateless policy gradient. Hence, for A 0, both bounds reduce the bounds found by prior work Papini et al. (2022) Zhao et al. 2011 when considering the univariate Gaussian case.
To understand the difference in the variance of both policy gradient estimators, need to compare AnPrT and As2GG. For S2PG, the additional variance is induced by TH2, which is the squared norm of the gradient of internal transition kernel rie (st, z4) w.r.t. 0, weighted by the ratio of the norms of the covariance matrices. In contrast, Theorem 3.44fiighligtt the effect of backpropagating gradients through a trajectory as ABPTT additionally depends on the constants 2 and Z. First of all, for H > 1, K >> and Z we need to distinguish two cases: as detailed in Lemma A.2 in Appendix [A.61]] for Z 2 1, the 2 Z grows exponentially with the length of the trajectories T, indicating that exploding gradients cause exploding variance, while, for Z < 1, the gradients grows linearly with T. That is, while the variance of BPTT heavily depends on the architecture used for (s, 2;), the variance of S2PG depends much less on the architecture of the internal transition kernel allowing the use of arbitrary functions for (7, 2;) at the potential cost of higher - yet not exploding variance. Hence, S2PG does not rely specialized architectures that account for exploding gradients such as, GRUs or LSTMs. We note that, for S2PG, we can define a tighter bound than the one in Theorem 3.5 under the assumption of diagonal covariance matrices shown in Lemmal A.3 in Appendix A.6.2
can observe that the quadratic term H 222 2" dominates ABTTT- Secondly,
EXPERIMENTS
We evaluate our method across four different types of tasks. These tasks include MuloCo Gym tasks with partial observability, both with and without a privileged critic, which we consider as Robust RI. experiments, complex II. tasks with a policy under partial observability and a privileged critic, memory task, and task in which we train the parameters of an ODE in a policy. Figure provides an overview. To conduct the evaluation, we compare recurrent versions of three popular RL algorithms-SAc. TD3, and PPO - that employ S2PG and BPTT. In Appendix C we present the algorithm boxes for all S2PG variants. Furthermore, we compare against stateless versions of these algorithms, which use window of the last 5 or 32 observations the input to the policy. To show the importance of a policy state, we include basic versions these algorithms, where only the observation given to the policy (vanilla) and where the full state is given the policy (oracle).

---

Published as a conference paper at ICLR 2024
Priviledged Information
No Priviledged Information
1.0
30
1.0
Norma lerdallzeeddee Return
00.8
Bog Banuned Time (h)
20-
0.6
0.6
15
Training
0.2
0.2
(Ours)
BPTT-5
BPTT-32
Window-5
Window-32
Vanilla
Oracle
AtlasPOMDP-Carm
HumanoidPOMDP-Walk
HumanoidPOMDP
Cumulative Reward
Cumulative Reward
Cumulative Reward
800
750
400
500-
400
250
200
500
1000
1500
2000
1000
1500
2000
5000
10000
15000
Steps (x10*)
Steps (x10²)
Steps (x10²)
GAIL-8PTT-S
LSIO-8PTT-5
GAIL-SS (Ours)
LSIQ-vanille
GAILBPTT-32
LSIQ-8PTT-32
LSIQ-RS (Ours)
GALL-vanila
Figure 3: Top [Results Robust RLJ: Comparison different RL agents on POMDP tasks using our stateful gradient estimator, BPTT with a truncation length of 5 and 32, the SLAC algorithm and stateless versions of the algorithms using window of observation with length and 32. The "Oracle" approach is a vanilla version of the algorithm using full-state information. Results show the mean and 95% confidence interval of the nor- malized return and the training time needed for million steps Imitation Learning]: Comparison of different versions LS-IQ and GAIL different imitation learning tasks under partial observability, Results show the mean and 95% confidence interval of the normalized return.
all POMDP Gym tasks. Bottom [Results
To compare with specialized algorithms for POMDPs, we compare with SLAC Lee et al. (2020a) SLAC model-based approach that learns complex latent variable models that used for the critic, while the policy takes a window observations as input. Agents that utilize our stochastic policy transition kernel denoted by the abbreviation "RS", which stands for Recurrent Stochastic. We use Gaussian policies - for a and 2 in S2PG for all our experiments. The network architectures are presented in Appendix [ The initial internal state always set O. When an algorithm has replay buffer, we update its internal states when sampling transitions to resemble BPTT. For fair comparison, all methods, except SLAC, are implemented in the same framework, MushroomRL D' D'Eramo et al. (2021). To properly evaluate the algorithms' training time, we allocate 4 cores on our cluster and average the results and training time over 10 seeds for all experiments.
Robust RI. Tasks. The first set of tasks includes the typical MuloCo Gym locomotion tasks to test our approach in the RL setting. As done in Ni et al. (2022). we create partial observability by hiding information, specifically the velocity, from the state space. Furthermore, each link (c.f., Fig. 3). As described in Section 2.4] is challenging to learn a stochastic policy state while using bootstrapping. Therefore, we
randomize the mass of
privileged critic with vanilla FFN for bootstrap- ping approaches. On the top-left of Figure 3 we show the cumulative returns and training times for TD3 and SAC-based agents. The overall results show that our approach has major computa- tional benefits w.r. BPTT with long histories at the price of a slight drop in asymptotic performance for SAC and TD3. We found that our approach performs better on digh-dimensional observation spaces e.g., Ant and Humanoid and worse on low- -ww-dimeniional ones. A more detailed discus- sion about this given in Appendix E.44 While the window approach has computational benefits compared to BPTT, it did not perform well with increased window length. The full experimental campaign, learning curves, and more detailed task descriptions are in Appendix E.4] In the second set of tasks, labeled as NN Privileged Information' tasks, we aim to showcase the effectiveness of our approach without using privileged information, combining the critic with Monte Carlo rollouts. Employing the same tasks in the privileged information setting but without mass randomization, our algorithm differs from BPTT by not utilizing a recurrent critic. The overall results, presented on the top-right of Figure 3 consistently demonstrate method's superior performance over BPTT.
8

---

Published as a conference paper at ICLR 2024
While the training time of PPO-BPTT with truncation length of 5 is similar to approach, notable differences emerge with longer truncation lengths. Both BPTT and the window approach exhibit worse performance with longer history. Comparisons with specialized methods, such as SLAAC Lee et al. (2000a), highlight our method's better performance and faster training. Further results, runtimes, and discussions are provided in Appendix E.3
Imitation Learning Tasks. The main results of our paper are presented at the bottom of Figure Here, we introduce two novel locomotion tasks. The first task is Atlas locomotion task, where a carry-weight is randomly sampled, but the weight is hidden from the policy. The second task Humanoid locomotion task under dynamics randomization. The goal is imitate a certain kinematic trajectory- either walking or running without observing the type of the humanoid. To generate the different humanoids, we randomly sample a scaling factor. Then the links are scaled linearly, the masses are scaled cubically, the inertias are scaled quintically, and the actuator torques are scaled cubically w.r.t. the scaling factor. For both tasks, the policies only observe the positions and the velocities of the joints. Forces are not observed. As can be seen in Figure 3 our approach can be easily extended to complex IL tasks using GAIL and LS-IQ. We observe that in the IL setting our approach is able to outperform the BPTT baselines even in terms of samples. All results are given in Figure 14 and Figure 15 in Appendix E.3 where we also further discuss the results.
Memory Task. To show that approach encode longer his- tory, we present two memory tasks. The first task is to move a point mass to a randomly sampled goal from randomly sampled initial state. The task is shown in Figure 2 While the goal is shown to the policy the beginning of the trajectory, is hidden when the point mass moves away from the initial state. As can be seen in Figure 4 our approach can outperform BPTT even for longer horizons. the second task, the positions of two doors in a maze are shown to the policy when close to the initial state and are hidden once going further away. Figure provides an example and the results. As can be seen, our approach outperforms BPTT on short horizons while being slightly weaker longer ones in the second task. We expect that the reason for the drop in performance is caused by the additional variance of our method in combination with the additional absorbing states in the environment. The additional variance in our policy leads to increased exploration, which increases the chances of touching the wall and reaching an absorbing state. Nonetheless, our approach has significantly shorter training time, analogously to the results in Figure All results and a more detailed discussion are given in Appendix E.1
0.8
0.6
0.0
500
1000
1500
Steps (x100)
SAC-BPTT:S SAC-BPTT-10 SAC-RS (Ours)
Figure 4: Point mass results.
ODE Task. Until now, all experiments were based RNNs as stateful policies. As mentioned be- fore, our gradient estimator could be used to train the parameters of a policy containing an arbitrary ODE. Such policy is shown in Figure This setting particularly interesting encode inductive biases directly into policy. To provide proof of concept, we successfully train a completely blind policy on the Gym HalfCheetah and Ant tasks, where we used the ODE of CPG to encode an oscillation directly into the policy; a bias that is very popular in locomotion research. The ODE is simulated using the Euler method. More results and discussions are provided Appendix E.2
CONCLUSIONS
This work introduced S2PG, an alternative approach for estimating the gradient of stateful policies without using BPTT by exploiting stochastic transition kernels. S2PG is easy to implement and computationally efficient. We provide a complete foundation theory of this novel estimator, allowing its implementation in statc-of-the-art deep RL methods, and conduct a thcoretical analysis on the variance of S2PG and BPTT. While this method still cannot replace BPTT in every setting, in the most challenging scenarios and in the IL settings, S2PG can considerably improve the performance and learning speed. Unfortunately, while S2PG can replace BPTT in the computation of the policy gradient, it is still not able to properly learn a value function and an internal policy state at the same time in the partially observable setting. While we show that our approach works well when the critic is estimated with Monte-Carlo rollouts, we limit ourselves to the setting where the critic has privileged information for bootstrapping methods. We plan to investigate solutions to this issue in the future. Finally, we will investigate how our approach scales using massively parallel environments.
9

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
Calculations for this research were conducted on high-performance computers Lichtenberg and CLAIX at the NHR Centers NHR4CES at TU Darmstadt and RWTH Aachen (project numbers p0020307 & p0021606). This work was supported by the German Science Foundation (DFG) under grant number SE1042/41-1. Research presented in this paper has been partially supported by the German Federal Ministry of Education and Research (BMBF) within the subproject "Modeling and exploration of the operational area, design of the AI assistance as well as legal aspects of the use of technology" of the collaborative KIARA project (grant no. 13N16274).
REFERENCES
Firas Al-Hafez and Jochen Steil. Redundancy resolution as action bias in policy search for robotic manipulation. In CoRL 2021
Firas Al-Hafez, Davide Tatco, Oleg Arenz, Guoping Zhao, and Jan Peters. LS-IQ: Implicit reward regularization for inverse reinforcement learning. In Proceeding of the International Conference on Learning Representations, Kigali, Rwanda, May 2023.
Bram Bakker. Reinforcement learning with long short-term memory. In Proceeding of the Fifteenth Conference on Neural Information Processing Systems, Vancouver, British Columbia, Canada, December 2001.
Jonathan Baxter and Peter I.. Bartlett. Infinite-horioon policy-gradiens estimation. Jounral of Artifi- cial Intelligence Research, 15, 2001.
Guillaume Bellegarda and Auke ljspeert. Cpg-rl: Learning central pattern generators for quadruped locomotion. IEEE Robotics and Automation Letters, 7(4):12547-1254, 2022a.
Guillaume Bellegarda and Auke Ijspeert. Visual cpg-rl: Learning central pattern generators for visually-guided quadruped navigation. arXiv preprint arXiv:2212.14400, 2022b.
Luigi Campanaro, Siddhant Gangapurwala, Daniele De Martini, Wolfgang Merkt, and loannis Havoutis. Cpg-actor: Reinforcement learning for central pattern generators. In Towards Au- tonomous Robotic Systems: 22nd Annual Conference, TAROS 2021, Lincoln, UK. September 8-10, 2021, Proceedings, pp. 25-35. Springer, 2021.
Kyunghyun Cho. Bart van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: mecoder-decoder approaches. In Proceedings of SSST-8, Eighth Workshop on Syntax, Semantics and Structure in Statistical Translatio, pp. 103-11, Doha, Qatar, October 2014.
Younggil Chol, Sajjad Manzoor, and Youngjin Choi. Adaptation to environmental change using reinforcement learning for robotic salamander. Intelligent Service Robotics, 12:209 2018, June 2019.
CoR1.2021 Proceeding the Conference Robot Learning, London, United Kingdom, Novem- ber 2021.
Carlo D' D'Eramo, Davide Tateo, Andrea Bonarini, Marcello Restelli, and Jan Peters. Mushroomrl: Simplifying reinforcement learning research. Journal of Machine Learning Research, 22:1-5, 2021.
Yonathan Efroni, Chi Jin, Akshay Krishnamurthy, and Sobhan Miryoosefi. Provable reinforcement learning with short-term memory. In ICML. 2022
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley. Iain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. ICML 2018
Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation enror in actor-critic methods. In [CML, 2018
10

---

Published as a conference paper at ICLR 2024
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- maximum entropy deep reinforcement learning with stochastic actor. In ICML 2018
Nicolas Heess, Jonathan J Hunt, Timothy P Lillicrap. and David Silver. Memory-based control with recurrent neural networks, arXiv preprint arXiv:1512.044 2015.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Proceeding of the Thirtieth Conference on Neural Information Processing Systems, Barcelona, Spain, December 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 9(8): 1735 1780, 1997.
(0182018. Proceeding the International Conference on Machine Learning. Stockholm, Sweden, July 2018.
ICML 2022. Proceeding of the International Conference Machine Leurning, Baltimore, Mary- land, July 2022.
Maximilian Igl, Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational reinforcement learning for pomdps. In [CML 2018
Auke Jan Ijspeert. Central pattern generators for locomotion control in animals and robots: a review. Neural networks, 21(4):642-655 653, 2008.
Auke Jan Ijspeert, Alessandro Crespi, Dimitri Ryczko, and Jean-Maric Cabelguen, From swimming to walking with a salamander robot driven by a spinal cord model. science, 315(5817):1416-1440) 2007.
Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning. In Proceeding of the International Conference on Machine Learning, Sydney, Australia, July 2002.
Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. Proceeding of the Thirty-fourth Conference on Neural Information Processing Systems, Virtual, December 2020a.
Jonathan N Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill. Supervised pretraining can learn in-context reinforcement learning. arXiv preprint arXiv:2306. 14892, 2023.
Joonho Lee, Jemin Hwangbo, Lorenz Wellhausen, Vladlen Koltun, and Marco Hutter. Learning quadrupedal locomotion over challenging terrain. Science robotics, 5(47)xeabc5986, 2020b.
Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In Pro- ceeding of the International Conference Learning Representattinss, San Juan, Puerto Rico, May 2016.
Puze Liu, Davide Tateo, Haitham Bou-Ammar, and Jan Peters. Robot reinforcement learning on the constraint manifold. In[CL 2021
Lingheng Meng, Rob Gorbet, and Dana Kulić, Memory-based deep reinforcement learning for pomdps. In Proceedings of the 2021 IEEE/RSJ International Conference Intelligent Robots and Systems, pp. 5619-5626, Virtual, September 2021.
Takahiro Miki, Joonho Lee, Jemin Hwangbo, Lorenz Wellhausen, Vladlen Koltun, and Marco Hut- Learning robust perceptive locomotion for quadrupedal robots in the wild. Science Robotics, 7(62):cabk2822, 2022.
Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov. Recurrent model-free rl can be strong baseline for many pomdps. In ICML 2022, pp. 16691-16723.
Matteo Papini, Matteo Pirotta, and Marcello Restelli. Smoothing policies and safe policy gradients. Machine Learning, (11(11):4081-111 2022.
11

---

Published as a conference paper at ICLR 2024
Xue Bin Peng, Marcin Andrychowicz, Wojeiech Zaremba, and Pieter Abbeel. Sim-to-real trans- fer of robotic control with dynamics randomization. In Proceeding of the IEEE International Conference on Robotics and Automation (ICRA), Brisbane, Australia, May 2018.
Lerrel Pinto, Marcin Andrychowicz, Peter Welinder, Wojciech Zaremba, and Pieter Abbeel. Asym- metric actor critic for image-based robot learning. In 14th Robotics: Science and Systems, RSS 2018. MIT Press Journals, 2018.
Ilija Radosavovic, Tete Xiao, Bike Zhang, Trevor Darrell, Jitendra Malik, and Koushil Sreenath. Learning humanoid locomotion with transformers. arXiv preprint arXiv:2303.0338) 2023.
Kate Rakelly, Aurick Zhou, Deirdre Quillen, Chelsea Finn, and Sergey Levine. Efficient off-policy meta-reinforcement learning via probabilistic context variables. In Proceeding of the Interna- fional Conference on Machine Learning, Long Beach, California, July 2019.
John Schulman, Nicolas Heess, Theophane Weber, and Picter Abbeel. Gradient estimation using stochastic computation graphs. In Proceeding of the Twenty-ninth Conference on Neural Infor- mation Processing Systems, Montreal, Canada, December 2015a.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In Proceeding the International Conference Machine Learning, Lille, France, July 2015b.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv, August 2017.
Haojie Shi, Zhou, Hongsheng Zeng. Fan Wang, Yueqiang Dong, Jiangyong Li, Kang Wang. Hao Tian, and Max Q.-H. Meng. Reinforcement learning with evolutionary trajectory generator: A general approach for quadrupedal locomotiong. IEEE Robotics and Automation Letters, 7(2): 3085 - 3092, April 2022.
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.
Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. In Proceeding tfe Thirteen Conference on Neural Information Processing Systems, volume 12, Denver, Colorado, USA. November 1999.
Daan Wierstra, Alexander Förster, Jan Peters, and Jürgen Schmidhuber. Recurrent policy gradients. Logic Journal of the IGPL, 66666460 2010.
Zhihan Yang and Hai Nguyen. Recurrent off-policy baselines for memory- based continuous control. In Deep Reinforcement Learning Workshop, NewrIPS, Sydney, Australia, December 2021.
Marvin Zhang, Zoe McCarthy, Chelsea Finn, Sergey Levine, and Pieter Abbeel. Learning deep neural network policies with continuous memory states. In Proceeding of the IEEE International Conference on Robotics and Automation (ICRA), pp. 520-527, Stockholm, Sweden, May 2016.
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.
12

---

Published as a conference paper at ICLR 2024
A PROOFS OF THEOREMS
Within this section, we refer to the following regularity conditions on the MDP taken from form Silver et al. (2014):
Regularity Conditions A.I. - (8, P(s'|s,a) a), VaP(s'ls, a), to z'|s, Vomola, 2'|s, z), (8, 2), Veue(s, =), r(s.a). a)) ((s) are continuous in all parameters and variables
8, a, 2, O.
Regularity Conditions A.2. There exists b and such that sup, c(s) < b, sups .. p()))s.)) <
b, su r(s,a) < b. supa Var(s' |s, a)ll < I,, and sup (ar( a) < L.
A.1 STATEFUL POLICY GRADIENT WITH THE SCORE FUNCTION
Proof. Let
T-1 J(F) = Σ a))(
be the discounted return of a trajectory 7 ($g. 20, a1.... ST 11- 2T 0r-1, 8T, 2T). Differently from the standard policy gradient, we define the trajectory as a path of both the environment and the policy state. The probability of a given trajectory can be written as
T-1
p(r 0) (80, 20)
II A=
P(144 (st, at ( 241 21).
The continuity of c(su. 20). P(st+1 st. a.), Ta(a,, 2411/5; Vo 2141/ 21) implies that p(F ()) and Ver(F(0) are continue well. Hence, we can apply the Leibniz rule to exchange the order of integration and derivative, allowing us follow the same steps the classical policy gradient
VoJ(To)= = Vo p(F())](F)dF= Vop(F(0))(T)dF
Applying the likelihood ratio trick, we obtain
Vo.T(TO)) = / p(Fl0)Vo (o) ([o
To obtain the REINFORCC-style estimator, we note that
log (() = log (son, za) + log P(stt az)
log no (ar 2441 21)
The gradient Vep( F 0) only depends on the policy and the initial state distribution or the state transitions. Therefore, we can write the stateful policy gradient as
Ve log (a 2411 22))(()
VoJ(TT) =
(10)
concluding the proof.
A.2 STATEFUL POLICY GRADIENT PARTIALLY OBSERVABLE ENVIRONMENTS
Proof. The derivation follows exactly the one for the fully observable scenario, with the difference that we need to deal with the observations instead of states. The probability of observing trajec- tory F in POMDP is
0 <(su, 20) (oulss))
II
O ((ovea/8t41 P( 1ls, ai 2411 10, (o,, z1).
Computing the term Vo log p ((1(0) and replacing it into equation we obtain
VoJ(TO) =E
Vo log me(os, 2041 lo,, x))(F)
using the same simplification as done in equation 10
concluding the proof.
13

---

Published as a conference paper at ICLR 2024
A.3 STATEFUL POLICY GRADIENT THEOREM
Proof. We can write the objective function as follows
To
E [40-20)>=
[V (50, zo)]
E
1 20, 0o, 21)]
(00. (00.21)~<pp)
We can now compute the gradient as
VoJo = Vo
(s0 20, ao, 21)
(20.20)-0 Jos.,
(ap loig
= Vo 1
20) =o: do, =1)dagdz1 dagdzo.
The regularity conditions A.1 imply that (s, a), and VaQ**(s, (s, a0,z z) are continuous functions of 0, s, a, and 2, Further, the compactness of S, A and Z implies that for any 0, [|VoQ* (s. z", a, =| and ( z'|s, 2)Il are bounded functions of s, and z. These conditions are required to exchange integration and derivatives, as well as the order of integration throughout the this proof. Hence, the gradient can be written as
VoJo [ [ Vo [c(so, zo)rr( 21|80, 20 =0) ( 20, ao, a1) daodz doodæ0
[ [ ((su, zu) VoTo 20/ ) (su; 20, au 21)
(80, 20) TT (Co, So. 2n) VoQ" (so, zo. do, 21 dagd dsoday..
We now focus the gradient of the Q-function w.r.. the policy parameters. Using the following relationship
( St, at, r(s., +
We write
(s,, = Vo (((s1, as) + P( (18==)or, (8141, 20 dd841
((s441)s,,,1)V4 (
Similarly, we can compute the gradient of the value function as
VoVTo 21) = [ale (Vonolus. 24t1 |s,, 5 (st, 21, ar, 21+1)
47 [ ma(ur, 241181, 2x) P(s441 V. TM V (1411, destidendze1..
We now focus again on VoJo. Using the linearity of integrals, we write
VeJo
((s0, z0) VoTo z1 (sp, zo QQ (so. zo, ao, z1 31)dadd dagdzo
(Ssu, 20) Zu ( 21 20) V T 20, an 21 )aadd daudia-
14

---

Published as a conference paper at ICLR 2024
Now, using the expressions for VoV and VeQ=*, expand the second term of the previous sum for one step (limited to one step highlight the structure)
(Bus 20 (au 2o su. 20) VoQ (Su, 2o, on, 21 )daodardsada
((80, 20 20 00 V 2odd
(50, 20 mo (a: 21
ao)
Vato 22 (81,
+y To(as, 21 21 P(82 81, an Vo Voveo (2))dd ds dooddaddsodag
By splitting the integrals again and rearranging the terms, obtain
1,5;; TO (ao. 21 s0, 20) VoQ (80, 20, oo, 21 jdgozzdsadda (ao:
[
J (88, (ao, 21 so, P(s so, ao daddddd
VAJ / Ve (a1, #2|s1, 21 2 (81, 21, 01, 22 22) dsi
s A 1 JaVel Ss (sso, ( 21 (a)) 22/11, ((s))s,))) Vel V no (82, 22 22)2ddd (ilisddyddzdss d dandz; d
To proceed and simplify the notation, we introduce the notion of t-step transition density Tr, i.e. the density function of transitioning to a given state and policy state in t steps under the stochastic policy me. In particular, we can write
Tre (50, 20) c(so, 20),
Tr" (81,
((su, zo)mo((a, 21|s0, 2os)a(((ssss.adddddddaoooo ao)dsadzzddao
Tr." (8,, #t)
((80, 20)
II
molas 2x+1 (sk, 2n P P(su+1 daaddaada.
KmD
Now put everything together, expanding the term containing Vo V infinite amount of times
[// Tr (80, 20)
" (4o, zo, ao, dadddaddddsdda Sale VoTo 21/s0, =0
VoJo =
/ / 20 (s), (5,, ar, dod dz
+
Tt (51,21)
+ 2
Tr 20 (82, 22)
Vona 23 (52, ويح au, )azzzdd
4
TT (83, 23)
Vono 2183, (83, 23, as, ddggdd
+...
[
Σ
Tr. (st,
VoTolan 21, 2x) 21, ar, Ze41 dda ds
15

---

Published as a conference paper at ICLR 2024
Exchanging the integral and the series obtain
VeJo
Σ
Tr?"" ze)-
t=
Sale Vomola, 2411511 20 (st, 2a, au 2411)aaa dssidz..
Now, with the definition of the occupancy metric
(8. 2)
T
and by relabelling St, 24, at 2141 into 8, 2, a and obtain
VoJo 161 (00(0,)) Lale Voso(a,z'|s,2) (4,,,,, d') da'ds da.
A.4 STATEFUL PERFORMANCE DIFFERENCE LEMMA
This lemma is particularly useful because allows using the advantage function A computed for an arbitrary policy to evaluate the performance of our parametric policy те. As is an arbitrary policy not depending on 0, the left-hand side of the Performance difference lemma can used as surrogate loss. This lemma allows to easily derive S2PG version of PPO and TRPO.
Lemma A.1. (Performance Difference Lemma for Stateful Policies) For any stateful policy ı and q in an arbitrary MDP, the difference of performance of the two policies in terms of expected discounted return can be computed as
[(]
Σ r ( 2e, Ge, 2011) (=0
Proof. Starting from the lemma
-(m) - = P
rewrite the advantage
( 21, at, Zt+1) = E (F(se, 0z) + V (o,, 21)
We can then manipulate the right-hand side as follows
E TNN.P
(st, 2e, an 2011) Σ
= E
r
( (1),
P
Σ Y'r(s;, ax)
E P
E
(+1 ( Σ In (s,
Twan P
-I(T) +
E
(V 21) - (s,
Jon.P
=I(T) + E
[V]
(0,2)
=J( -T(9),
concluding the proof.
16

---

Published as a conference paper at ICLR 2024
A.5 STATEFUL DETERMINISTIC POLICY GRADIENT THEOREM
Proof. The proof proceeds in a very similar fashion as the one for the standard DPGT. We start by computing the gradient of value function arbitrary state and an arbitrary policy state
VaVio (8,z) =VQQ (s. p0 (o,(;
=Vo ((s. P( (s, 2)) (s',, 2))dss
-Vouo(s, V + Vo / yP(''ss, 1 2H) Hā(s, 2))ss'
_oo1ā(s.)) (onnl(,.//arrs.)aaaaa00
(P(d v, Vov (o', nā z))
Vouā(>.:) ( M(s. 1s, ds'
=VoHōls22 r(s.a)+ + / TP 1s, a)) ()(ss (
7P(5'ss
+ vouā(s,2) ( () V. V 2 dds'
| p((d)s, VoVHo (6', 2') mg(0.)) do'.
(11)
Now, we can observe that
(s, V. (r(s.a)+ (5
= = P(s''s, (d, d
(12)
Substituting equation into equation
11 yields
VaV*o (s =Voue(s, V). (s, uā 2)) ==0=
VoHb(s,2) V. 18( 2))
| P(s''s, VoV"M (s' =²) = do'
(13)
To simplify the notation we write
Ge (s,2) = Voro(s, 2,, :)(x, azog)
=xgg()))))) VoHb(s,2) V T. (8 2), 2))
Therefore, equation 13 can be written compactly as
VoV" = =e 2) +
7P( -( (s (2),,")))==gg00 do'
(14)
Using equation 14, we can write the gradient of the objective function as
VoJ(pp) =Vo ((s0, V (so: zn dspddo
= [ / ((su, 20) Gte (su, 20 dsaddo
[ [ (800, (000)20 20) /
(000 20)) Vo (11, (i)l dsıdsydao. (15)
Before continuing the proof, we need to notice that the variables s' and of are conditionally indepen- dent if the previous states and are given. Furthermore, it is important notice that, even if the
17

---

Published as a conference paper at ICLR 2024
policy state is a deterministic function of the previous state and previous policy state, the variable 2 is still a random variable if we consider its value after more than two steps of the environment, as its value depends
set of random variables, namely the states encountered during the path. Using these two assumptions, we extract the joint occupancy measure of s and 2. Unfortunately, to proceed with the proof we need to abuse the notation (as commonly done in engineering) and use the Dirac's delta distribution ((r).
Due conditional independence, we can write the joint transition probability under the deterministic policy μig as
pro (s' [s, puo (s'|s, 2)8"0 (z'1x, 2) p(s z))8(()))))) no(s,))),
with ppo (s'[s; 2) = P(s'|s, (s, :)) and 8 (2'1s, - ó(=' - 1 (s)))))
Now we can write
/ on (s', 2s1,, = on (s,1s,))0000 (z' 2)d2'.
pse (d'(s,2)
We introduce, as done for the PGT, the t-steps transition density Tr." under the deterministic policy Ho
( 20) = 1(50, 20),
PHe - ((80, =0) pHe (s1lso, =0)000 (=1 ddoodoo;
Tr (51, 21) = 21)=
1-1
Tr/20 Tr/o (st; |Six2 1880, 2u)
II
pro (se+1 (1641//k. 20) (ze+1 lsk, zn ddsydu.
k=0
Using the previously introduced notation, expand equation 15 by recursion, infinitely many times
So ((su, ((8u;20) 20) (50) 20) + c(su, 2u) I yPu ($1, 21|8u; 20 VoV (81, 21 d1)d1dd1
dsudzo
The (5) (so, 20) dsgdzo + / [ ye(8o, 20) pro (81, 21 so, 20)
(01))) + 1 ppu (#2, (81, 21) VoV'0 (s2, 22 )sgddz dapiap
Js Tr (50, 20) douddo + "Tr (81, 21) (81, 21 ddai
PHe * [S3xz0 22((8. PHo (81, 31 ss, 20) pue (82, 22|31, Vo (82, (2) dd)
Σ E/
[ 7 "I (st, 24 Gn0 (st, (st, 2t)dsidzi.
(16)
Now, we notice that the occupancy measure is defined as
pda (x.2)
T'Tff (st, 21)
Therefore, by exchanging the order of series and integration we obtain:
VoJ(Ho) [ "Tr"" (s, =)(" z)dsdz
[ (s, (Voro(s, V. Q" (8, Hō(s, ())( 2))oo,
+Veul(s,2) V. Q"0 Hō(s 2), x²)) )
which concludes the proof.
18

---

Published as a conference paper at ICLR 2024
A.6 VARIANCE ANALYSIS OF POLICY GRADIENT ESTIMATES
the following section, we define the variance a random vector A (AI...., Ar) as in Papini et al. (2022); Zhao et al. (2011)
Var[4] =r (B (4 1EA4-I(AAA [AA]A E[AD"] )
(17)
Σ E (Am [[Am))
(18)
Using this definition, we can compute the variance of the policy gradient estimators for stateful policies using BPTT and S2PG.
A.6.1 BACKPROPAGATION THROUGH TIME
Proof. We consider policy of the following form
1 2 pa(he)))
(((lu)))===< 2- uo((h)))
plas/hr, @) = -aaa (k)))
exp
(19)
2xd
where ④ 10, }. For a fixed covariance Σ, REINFORCE-style gradient estimator for this policy depends the following gradient
f(r)
Σ
Ve log T(ache,, 0) Σ
Vopo(h))Vn lag Mmmp m
Σ Vopo((ee (E - -ma) Immunothe)
Given the total discounted return of a trajectory
G(r) t a(1, 8141),
the variance of the empirical gradient approximator given by
Var [oJ(o)]
1 Var (G(T)
(20)
N
where is the number of trajectories used for the empirical gradient estimator. Therefore, we can just focus on Var (G)) Using equation 18 we can define an upper bound on Var [G(T)
Var [G(T)
[GG]]]
= p( p(r)
Σ 1-1 r
Vapa(lu) (I- (a. -ma)) mi=
Vauo(hu) (I- (as mi)) (mi=po(te)
Let & I- mi) where N(m,, E), then (((m)) are random variables drawn
from a Gaussian distribution (O, E-I). Hence, we define p(E.) E-, and treat & as an independent random variable sampled from p(E.). For compactness of notation, define [r) where p(E) P(E.) In similar way define
19

---

Published as a conference paper at ICLR 2024
p(T(E) ((50) Io P(s Et). Using this notation, can write
Var (G(T)[(T)] p(T(E)
rss, des 8t-1)
Vopsa(hee)s
Σ (apa Ex)
R² - -²
E
- 9)-
7.6
(1
I3
& Vopo(h)) Vopo(he)
(21)
(1 7)
Ti
+ riC
& (opu(lu) Vapa(be)) Es 1.0-0,
Σ
Now, we take a look at the last term. Without loss of generality, we will assume > t. As & is random variable drawn from a Gaussian distribution with zero mean, and as he and he are independent from fr, we can write
E
Σ
& (oosoo(he) Vopo(le)Ee
T-1
Σ
E
Ee (apishh) )omo(he))e))
1,4 1a
E Voue (ope((he) B [Ev] Esl
Ee
7.
1,1=0
0.
where in the last line used the fact that Eg,, [Ev]] = -. The same result holds for > t, following a similar derivation.
Continuing from equation 21 let E, = (hu T Vola
Noticing that E, is square positive
semi-definite matrix, we can write
22 A
Σ
Vouo(he) (osl(h))EE
7)2
T-1
E
E Vopro(he) [opp((n)e
(1- 7)
R²(1 (111-1)² T-1
Σ
(r (Ho)) (opo(k))==))
7)
Tit
R² PP1-77P
[optettt)) I| 7.6
(oaa((lalle ||=- [le]
(1 (1-7)² 7)
(1 713
E (
(1 7)
where I1-IIr is the Frobenius norm. When μ(he)) is implemented recursive function j(hh)) fo(s,, 1)-1) where 2,+1 (s 24) the gradient is given by
8 Vouu(hy)= 50 fo 2.)
Σ 1=0
no (s1. 23)
II 1=
Q Ozy
no (5))
0 dz
fo(
(22)
30
20

---

Published as a conference paper at ICLR 2024
Using this gradient and using the Assumptions 3.2 and po(ht) follows
3.3
we bound the norm of the gradient of
Vapao
-
20
+
Σ
a0
0
(s
zi)
1I
8
(* 25)
111))
| fo 24)
F
#==+1
S
Σ
K
j=
= H
z
K.
Hence, the final bound given by
R² | || 72
2
Var [oJSPI((O)] -
Σ
F Σz [))
N(1 - 7))2
R$
[| 7²
(T + +2HKK #
N(1 772
where define 7 = Et= Eiz ziIt1 and 7 Σi E/ - cluding the proof.
2
for brevity,
We can further analyze the properties of the BPTT estimator with the following lemma:
Lemma A.2. For Z < 1 the two constants 2 and 2 grow linearly with T, while for Z > 1 the two constants grow exponentially with T.
Proof. The series expressing Z can be written
2
Σ
Σ
T-1
Σ
(-1
T-1 Σ
(1
ZY T-1 Σ z
[m
k=
1w0
T-1 T Σ
2
(1- (1-Z)²
1=0
1=0 1-2T 1-Z
T
1- 2
+
(1 2)²
1- 2e
T
1- ZT - 2J (11 2)
(1 z)²
2)²
(1-2)(1+2 Z)
T
(1 - z
1(1 2TT - 2 2)
(1 - 2)² +
(1- (1= 2)
T
(1 2- (Z= 22 1))
+
(1 - Z)2
(1 (1-Z)3(12))
T
1Z -1)(2" 22 1)
+
(23)
(1 (1-2)3
(Z - 11+ 2)
21

---

Published as a conference paper at ICLR 2024
From equation one is O(Z2T)
23 we that Z is composed of two terms, where the first is O(T) while the second The series expressing % can be written
T-1 7
T-16 -
Σ 1-7
1-
t=0 1=0
== k=0
T
Z
+
(24)
1-Z 7
2
(1
2)
From equation 24 we see that 2 is composed two terms, where the first O(T) while the second one is O(Z"). These observations conclude the proof.
A.6.2 STOCHASTIC STATEFUL. POLICIES
Proof. For stochastic stateful policies, we consider similar policies of the form
Taa, 20t1/86; 2t, θ) ((00()).2441).1)
(25)
exp - (ar 0o(s1,z))) Tr) - -00(s,,2)))
1
2
where Ô {0, 11,, where 0, is the mean vector containing the means corresponding an action and hidden state such that do(s,, =x) [ (st, =,) (8 2.) T] and I is covariance matrix
Σ 0
containing the covariances corresponding to action and hidden state such that r
0
Then, we can write the gradient with respect to trajectory as follows
f(F)
#411 2, Ô) Σ Vo (logT(aelse, 20, Q) log (ogg((2411855) 2t, 6))
Vona 2x) Mi log (mme: Σ) ml moo(nne))
VoHă( 2.) V., (b, T)
Vong (st, =x) (x- me)) lou-roten))
VoHō(se, z)) pr-'(zee) - be)) (o:=p((on))))
where we have used the fact that and 2+11 are independent for decomposing the gradient into part corresponding to the mean of the Gaussian used for sampling action and part corresponding to the mean of the Gaussian used to sample the next hidden state.
Let & (-+ - by) where 2+11 Nhb,, Y), then (b:) are random variables drawn from Gaussian distribution (00 r-1). Hence, we define p(E.) = (0 T-1) and treat & as indepen- dent random variable sampled from p(E.). Then
equation 18 again to define upper
22

---

Published as a conference paper at ICLR 2024
bound on the variance such that
((F))(+) - [ p(E))())p(E
Σ x as,
Vepu 2r) Ee Vo 20 )
(8, 21) E Voru( 5
R² - )ー (1 - 7)2
E
Vonu 21) E + (0)
FiE.E
(Vopū(Se; Eer + Vopb(se,2))65
(1 (F)2-
Vonū((0, 2x VoNb(Hr,
(1 -7)3 TiciE
(Vopy =e) & + (()
Σ (VoH 2OHō(s(.21)..
E F.E.
(Vopō(se:2))E) + Vonū(se,2))E))
Again, the last term can be dropped due to
T-1 Σ
(Voré(a))))))
21) + Vopō(or, =)):) (Vorg(eu, w)) + Fonb(nn.w)E)))
T,EK
P=
(Vonsler + Vorālon, (Vonā(eu av) + VoHō(Su,
E
T.E.
(Vous(ee 21) VeHb(sn, ]
Σ
E
T.E
(Vonō(sr,2e) + Vopā(se,
E
F.g
P-1
Vonl 2)) 5o
Voni(oes 2x) Et
T.E.
r.
e, =D
R Vong (eu Ev
+ E T.
Vonā(su, 20)6r 0.
Noticing that (st Vopg(s 2) and Voug(s T Voug($ 24) are square positive semi-definite matrices, allows write
23

---

Published as a conference paper at ICLR 2024
A² -2 E
2e) + 2o) T 2c) Vo
(1 11-1)² 7.6
(1 )²
Σ
T-1
E ribiE
(Vene((s) + 2e) (atp://aaaatttaaaaaa/aaa/sstts-. (e
(1 7--
1=0
- {)
7)
(Vopō(e))2)))) (Von((n)))))))) + E (Vorg(on (Vouō(en Pi
+ te.K
(Vorū(e) Vo
+ R T,sK
(Vonū(Fe,2))
IP - )²
T-1
Σ tr (opt((ee;2)) 20) Vorb((
7)²
ir (Vonū(s),2)) 2)) Vopp(te, 2x)
$ A1 (1-1- (1 7)
Vorū x) [|= + [ Von -11-111--111 (|T=)11])
To align with the notation from the previous proof, define pg(s), 24) MD (8t, 21) = T/ (st, zz) and get
fo 21) and
-1 (1 7)
Σ
Volo
21)
-|= ||=
2x)))11 pille-11r-111 ||r-"1lee)
1=0
(1-17") 7)-
Σ F
|| H²
1|x-")|
such that the final bound on the variance of the S2PG is given by
|
R²
1| (1---)²
T-1
r 11
Var [oossvo(()]
Σ L
F H²
N(1 7)
concluding the proof.
Lemma A.3. Given a policy, similarity to equation 25
π (ar, Sst1/ 2i, θ) = (No 2+ 1)
Σ
with I
and limiting the covariance matrices Σ and Y to be diagonal, then the bound
0
T
from Theorem 3.5 be replaced by a tighter bound
R² tr(E-)111))) [11 PPT
F2 + H3 tr(r-1)
Var VoJssra(0))
7)1
ir(2-11)
Proof. Indeed, when is diagonal it is easy to show that
tr(AB) = diag(A) diag(B).
(26)
Let diag(E) [0o..... Tol T, and diag(T) can write
[vo..... U| T. Using equation 26
Vonu(ss, 21) Voui(es, 21) ) diag (Vope(ss, 2)) Voll 21) diag (E-1) (o1 VeHuloe, Vapulli Vep diag (VeHō((e; T (o) diag (I-').
24

---

Published as a conference paper at ICLR 2024
Let xll be the i-th row of the matrix X. We observe that
Voluū($6; 23) 2 ( Vo110 21)
diag (Vonō(ee, 2) (oo1((s, =
21) )1d (5 21) (all-11)
[V 24) 1
(s,, 11)1
The same observation can be done by exchanging 1g with pzg Notice that the resulting vector is vector of the norms of the gradient of the mean functions w.r.t a specific action (or internal state component).
If we assume the existence of two positive constants Fa and Ha such that ( 2 1113 - P2 and Veua 21 11 < Ha: Vi, St, 24 then we can write
tr (on (x))((1,2)))) E F2. (1 diag(E-1))
r l-
F22 tr(x
(Voug(sn 2)) Vopā(st, (e)ni(ss x
Ha (1" diag (Y (Y-
H H tr(
(27)
Using equation 27 the bound of the variance of the single gradient estimator is
11(1-17)
Var [M){{) <
Σ tr (Vonā 2)) Vo)10(0 x-1))
(1 7)
ir (VoHō(se, 21 )]
+ p
11(1--77
T-1
Σ
(F) (rr( + H1 t)(
(1 71-
R2 7TTTT
(F2tr(E))) tr(+)))
(1 7)²
Therefore, the gradient of the S2PG estimator in the diagonal Gaussian setting is
R R² TTT
(r(
Var [a)s
F H
N(1 7)2
ir(--1)
concluding the proof.
Notice that, when the action is a scalar, we get
Var [v)ssrr(@e]
<
(1 N(1 - -0) (F) + H trrttT ")==))
matching closely both the previous results of Papini et al. 2022 with the addition of the term depending on the internal state variance and the special scalar action case of the generic bound derived with a full covariance matrix.
25

---

Published as a conference paper at ICLR 2024
B QUALITATIVE COMPARISON OF THE GRADIENT ESTIMATORS
In this section, we analyze the difference between the stochastic gradient estimator for stateful poli- cies and the BPTT approach in greater detail.
In BPTT, the gradient needs to be propagated back to the initial state. This requires the algorithm to store the history of execution up the current timestep, if we want to compute an unbiased estimate of the policy gradient. The common approach in the literature is to truncate the gradient propagation for fixed history length. This approach particularly well-suited if the recurrent policy does not require to remember long-term information.
Instead, our approach incorporates the learning of the policy state transitions into the value function, compensating for the increased variance of the policy gradient estimate (coming from the stochas- ticity of the policy state kemnel) and allowing for the parallelization of the gradient computation. As shown in Figure [ the stochastic gradient for stateful policies only uses local information at timestep / to perform the update, while the BPTT uses the full history until the starting state 80- The choice between the two estimators is non-trivial: high dimensionality of the policy state may cause the Q-function estimation problem challenging. However, in practical scenarios, tasks can be solved with relatively low-dimensional policy state vector. In contrast, long trajectory implies multiple applications of the chain rule, that may consequently produce exploding or vanishing gra- dients, causing issues during the learning. This problem not present when using stochastic stateful policies.
0
81
21-1
43
81
81
40
5g
Vovolar,1le Q(st,
E Ez
(V (1, 2141 (81, 21) Q(si, 21, at 21+1)]
Figure 5: Comparison of the gradient at the state and action in BPTT -left and stochastic gradient estimator right- using PGT. The and green colors represents the value of future state, where green means high Q-value and red means low Q-value.
26

---

Published as a conference paper at ICLR 2024
C ALGORITHM PSEUDOCODE
In the following, we present the pseudocode for SAC-RS, TD3-RS and PPO-RS algorithms. The algorithms are straightforward modification of the vanilla SAC, TD3, and PPO approaches to sup- port the internal policy state. In all cases, the algorithm receives in input a dataset of interaction (in our experiments one single step for SAC-RS and TD3-RS, while for PPO-RS is a batch of experi- ence composed of multiple trajectories) coming from the current policy rollouts on the environment. In the pseudocode below we denote the parameter of the policy as 0, the target policy parameters O. Furthermore, we name the parameters of the i-th value function as vb and the corresponding target network 2h,. In the following, a vector. In this context, we refer to the vector of states, actions, and internal states inside the dataset (seen an ordered list).
the square bracket notation to denote the component of
Algorithm TD3-RS
1: function UP3ATETD3(D) M.ADD(D)
2:
b Add the dataset of transitions D to the replay memory M Wait until the replay memory has Smin transitions
3:
if SIZR(M) Smin then
4:
Sample a minibatch B from M (==COMUUTETARGET())
5:
6:
FIT(B, 2) for each i E {0,1} if mod 0 then Compute the loss
7:
D Perform the policy update every D iterations
8:
LB( = N
9:
< OPTIMIZE (B) 0)
10:
end if
11:
0+ 7 + (1 - (1- 7)
12:
bb 1 Tbb; + (1 T)
13:
+ +
Update the iteration counter
14:
end
15: return 0
16:
end function
17: function COMPUTETARGET(3)
for (s, a,,,, s¹) B do
19:
if s' is absorbing then
Ooeu(s' ((8,,'))
21:
else
22:
a',²' - 1o(6'' ; =²)
23:
Sample Ea, E- N/0, oz)
24:
<cpp
" + CLIP(E", 8e, de)
25:
Eap + CLIP(EE, dE)
26:
asmt
+ a' + Cdl
27:
Zemt
z" + E
28:
Oclp CLIP((ssmts amnn amax)
29:
Zchp CCIP (2amtt Zmin 2max)
30:
Unext (8' z/) +- min; Q (8', 2., Oclp; Zolp)
31:
end
32:
0(s, 2, a, 2') x- r Vanea (8',2')
33:
end for
34:
return أ
35:
end function
27

---

Published as a conference paper at ICLR 2024
Algorithm SAC-RS
1: function UPDATESAC(D)
2:
M.ADD(D)
3:
if SIZE( M) Smin then
4:
Sample minibatch B from M
5:
2 if SiZE((M) > Swarm then
6:
b Perform the policy update after Swarm samples
7:
Sample a' 2 To(-|s, 2) V(s, 2) e B
8:
Compute the policy loss
Lo = a" log ( [8, + a² log TogT((z' (=(1),,
1 Σ min (
(8, 2s,2,'',=')
N
9:
f- OP (B). 0)
Compute the a and a² losses with target entropies Ha and Hi=
Laa (B)
Σ a" (log (5(a), H()
N
(,,2,))EB
La (B)
Σ
(logT)(2'(s, R2-)
N
(5.5.2")
11:
a² + OPTIMIZE Lad (B)
12:
a² OPTIMIZE T... ())
13:
end
CCMMPTTSSOFTA Q..FIT(B, ②) 1/ Teb; + (1-r)ub
15:
16:
end
18: return 0
19:
end function
20: function COMPUTESOFTIARGEE(())
21:
for 2',,'','', B do
22:
if is absorbing then
23:
Unexi (8', zz)
24:
else
Sample a',²" (T(- Abonus + aa" log To (a' 1s,, 20)
26:
(2)====== a log (z(2"ss's
Unext 2/) t- min; Qú, (8', 2,, a, 2") + hbonus
28:
end if
29:
ğ(s, 2²) +- r W (s',2")
end for
31: return 9
32: end function
28

---

Published as a conference paper at ICLR 2024
Algorithm PPO-RS
1: function UPDATEPPO(E, D)
2: V A +-OMPUTEGAE(VV D)
3: FIT,FI((D, V
Update the value function
4: for i + to N do
5:
Split the dataset D in K minibatches [Ba k E [0, 1]} for k: to K do compute the surrogate loss on the minibatch
6:
7:
min 2'1s,,
Lo(Ba)
Σ
=(,,2 A(a, 2,,,),
N
(B,2.0,2")EBx
Tola, 2))
g/a 2(0,2155,2) 1 -eel, + A(s,2,
clip
8:
4 (B)) 0)
9: end for
10: end for
11: return 0
12:
end function
13: function COMPUTEGAE(VH, D)
for +- 0.. len(D) do
15: u[k] + Vo (s(k), [k])
Unext
[k] (s [k))
17:
end for
for
kiev <- 0. len (D) do
19:
l 4- len (D) Kics
20:
[k]] is last then
21:
if s [K]] is absorbing then
22:
A(s[k], 2(k) af + r[k] -[[k]
23:
else
24:
A(s[k), = alk:), (A)) + r[k] + 7Vneee [k] - v[ks]
25:
end if
26:
else
8 A( a + 1]. 2' 1])
28:
[ (k+)) <- r [A:] 7omee [k]] 2[k]] + 7/8
end
30:
end for
31:
return A
32:
end function
29

---

Published as a conference paper at ICLR 2024
D NETWORK STRUCTURES
In this section, we describe the structure of the networks used in this paper. Figure 6 presents networks used for the policies and the crities in the RL setting, For the BPTT, the policy networks were used, yet the red paths were dropped. Also, the critics did not use the hidden states as inputs. For the IL setting, the networks from SAC are used for LS-IQ, and the networks from PPO are used for GAIL. For the recurrent networks, we generally use GRUs even though any other recurrent network, such LSTMs, can also be used.
presents all
TD3 and PPO Policy
SAC Policy
Observation Embeeiee
fo
fo
Embeddier
Observation
Embeöder
RNN
Output Layer
Mau
RNN
log
Ri-
Emberide
Erobydder
7/0
22
7/9
X4
TD3 and SAC Critic
PPO Critic
State/Action Embedder
0g
MLP
MLP
V
Reo1
Embedder
Figure 6: Network architectures used for the policies and critics in PPO-RS, TD3-RS, SAC-RS, PPO-BPTI, TD3-BPTT, and BAC-BPTT. For the BPTT variants, the red paths do not exist. For the S2PG approach, the internal state is sampled from the noisy distribution. These architectures of the policies were originally used in Ni et al. 2022
30

---

Published as a conference paper at ICLR 2024
Randomly selected Goals
1.0
Figure Point mass experiments. In the blind area, the goal is not visible. The task involves remem- bering its position. Initial states, goals, and blind areas are randomly sampled. SAC-RS a recurrent- stochastic version of SAC, using S2PG.
Blind Area
0.6
0.4
Observable Area
0.2
0.0
1000 Steps (x10*)
Randomly selected initial Positions
SAC-BPTT-5
SAC-SPTT-10
SAC-RS
Goal
1000
Reward
Figure 8: Point door exper- iments. the blind area, the po- sitions of the two doors are not visible. The task involves re- membering its position, which randomly sampled. SAC-RS recurren-stochaste version of SAC, using S2PG.
600
Randomly selected Door Positions
400
200
Blind
O observable
0
1000
1500
2000
Area
Steps (x10*)
SAC-BPTT-5
SAC-BPTT-10
SAC-RS
E
ADDITIONAL EXPERIMENTS
E. MEMORY TASK
These tasks are about moving a point mass in 2D environment. They are shown in Figure 7 and For the first task, random initial state and goal are sampled at the beginning of each episode. The agent reads the desired goal information while it is close to the starting position. Once the agent is sufficiently far from the starting position, the information about the goal position is zeroed- out. The task consists of learning policy that remembers the target goal even after leaving the starting area. As distance to the goal involves multiple steps, long-term memory capabilities are required. Similarly, the second task hides the position of two randomly sampled doors in maze once the agent leaves the observable area. Once the agents touches the black wall, an absorbing state is reached, and the environment is reset. We test our approach in the privileged information setting using SAC. We compare our approach with SAC-BPTT with a truncation length and 10. For the first task, our approach- SAC-RS outperforms BPTT, while it performs slightly weaker compared to BPTT with a truncation length of 10 on the second task. We expect that the reason for the drop in performance is caused by the additional variance of our method in combination with the additional absorbing states in the environment. The additional variance in our policy leads to increased exploration, which increases the chances of touching the wall and reaching an absorbing state. Nonetheless, our approach has significantly shorter training time, analogously to the results shown in Figure 3 for both tasks.
&
E.2 STATEFUL POLICIES AS INDUCTIVE BIASES
Adding inductive biases into the policy is common way to impose prior knowledge into the policy Ijspeert et al. (2007); Bellegarda & Ijspeert (2022a); Al-Hafez Steil (202I); Liu et al. (2021). To show that stateful policies are of general interest in RL not only for POMDPs we show in this section that they can be used to elegantly impose an arbitrary inductive bias directly into the policy. In contrast to commonly used biases, approach allows learning the parameters of the inductive bias as well. We conduct experiments with a policy that includes oscillators - coupled Central Pattern Generators (CPG)- to impose an oscillation into the policy. The latter is common
31

---

Published as a conference paper at ICLR 2024
Ant-v3
Cumulative Reward
1400
Coupled Oscillators
Mymn y M May yo
1000
800
600
Z
500
1000
1500
2000
MLP
Steps (x107)
HalfCheetah-v3
14
2000
Cumulative
1000
TD3-R5 with CPG (Ours)
1000 1500 Steps (x10²)
2000
Figure 9: Experiments in which set of coupled central pattern generators are used as a policy. The left side shows the cumulative reward on HalfCheetah and Ant. The right side shows the structure of the policy. Note that the policy does use the environment state making blind.
for locomotion tasks Miki et al. 2002 We model the oscillators as a set of ordinary differential equations and simulate the latter using the explicit Euler integrator. Figure presents the results. Using our algorithm, we are able to train the parameters of the CPGs and an the additional network to learn simple locomotion skills on HalfCheetah and Ant. These policies solely rely on their internal state and do not use the environment state making them blind. With this simple experiment, we want to show that we can impose arbitrary dynamics on the policy and train the latter using our approach. Note that plain oscillators do not constitute a very good inductive bias. More sophisticated dynamical models that also take the environment state into account are needed to achieve better results.
Related Work. Shi et al. (2022) train CPGs using black-box optimizer instead BPTT, while simultaneously learning RL policy for residual control. Chol et (2019) optimize CPG policy using RL without BPTT by exposing the internal state to the environment but do not provide a theoretical derivation of the resulting policy gradient. [Campanaro al. (202T) train CPG policy using PPO without BPTT. However, the authors neglect the influence of the internal state on the value function.
E.3 COMPARISON ON STANDARD POMDP BENCHMARKS
In the following, evaluate the two gradient estimators in standard MuJoCo POMDP benchmarks. in the on-policy setting. These tasks are taken from Ni et al. (2022) and hide the velocity. All experiments the PPO algorithm approach outperforms BPTT in all benchmarks terms of the number of used samples and time. Unfortunately, using on-policy approaches these benchmarks, we do achieve satisfactory performances in all environments (e.g. Hopper and Walker). This is probably due to the increased difficulty of exploring with partial observation, rather than issue in gradient estimation. Indeed, in the other tasks, where more information is available, or the exploration is less problematic e.g., the failure state is more difficult reach we outperform the baseline. It should be noted that the performance gain in terms of computation time is not massive: this is due to the reduced number of gradient computations the on-policy scenario, which performed in batch after many environment steps. Table [ presents the training time needed to run million environment steps for all versions PPO. Accordingly, table illustrates the training time factors each PPO version in comparison to the vanilla variant.
learning method. Figure 10 presents the results. Our
32

---

Published as a conference paper at ICLR 2024
AntPOMDP-v3
HalfCheetahPOMDPPv3
HumanoidPOMDP-v2
Cumulative Roward
Namulaativ Reward
Cumulaative Reward
6000
4000
4000
3000
2000
1000
1000
3000
1000
3000
Steps (x10f)
Steps (x10*)
Steps (x10F)
HopperPOMDP-v3
Walker2dPOMDP.v3
Cumulative Reward
Cumullative Reward
2000
4000
1000
2000
1000
2000
3000
4000
1000
2000
3000
4000
Steps (x10")
Steps (x107)
PPO-8PTT-
PR0-6PTT-32
PPO-Vanilla
PRC-Wirdow-5
PPO-Window-32
Figure 10: Comparison of PPO with the proposed stochastic policy state kernel and PPO with BPTT on fully partially observable MuJoCo Gym Tasks. The number behind the BPTT implementations indicates the trunca- tion length. Abscissa shows the cumulative reward. Ordinate shows the number of training steps (x 10
33

---

Published as a conference paper at ICLR 2024
Figure 11: Example images of the Mujoco Gym randomized mass environment used within this work. The upper row shows three randomly sampled Ant environments, while the lower row show three Humanoid envi- ronments.
E.4 COMPARISON ON ROBUSTNESS BENCIMARKS WITH PRIVILEGE INFORMATION
In this section, we evaluate the performance of S2PG against BPTT on the POMDP setting with additional randomized masses. For the latter, masses of each link in model are randomly sampled +300 range. I examples of randomly sampled Ant and Humanoid environments. Here, we consider the scenario where the critic has privileged information (knowledge of velocities and mass distributions). while the policy uses only partial observability.
Figure 12 presents the reward plots. Figure 3 also presents the average runtime of each approach used for these tasks. The experiment run for maximum two million steps and four days of computation. The results do not show approach clearly outperforming the others in all tasks. S2PG based approaches seem to work more robustly high-dimensional tasks, such as AntPOMDP-v3 and HumanoidPOMDDP-2 learn in the AntPOMDP-v3 task. In general, S2PG seems to struggle in the HalfChettahPOMDDP HopperPOMDP-v3 and in WalkerPOMDP-x3 tasks. To better compare the performance of all approaches on the low-dimensional tasks - Hopper, Walker and HalfCheetah and the high- dimensional tasks Humanoid and Ant -, we presents the averaged performance plots Figure 13 As can be seen, our approach performs significantly better on the high-dimensional tasks while being worse on the low-dimensional ones. We believe that the reason for the worse performance of BPTT on high-dimensional tasks could be the potential risk of exploding gradients, which in turn cause exploding variance (c.f., Section 3). In contrast, we believe that worse performance of method diw-dimensional tasks can be traced back to the increased amount of variance in the policy in conjunction with the presence absorbing states in the case of Hopper and Walker, similar to the memory tasks. Table [ presents the training time needed to run I million environment steps for all versions of SAC and TD3. Accordingly, table illustrates the training time factors for each version in comparison to the vanilla variant.
Curiously, the BPTT version of TD3 not able
34

---

Published as a conference paper at ICLR 2024
AntPOMDP-v3
HumanoidPOMDP-v3
Walker2dPOMDP-v3
Cumulative Reward
4000
Cumulative Reward
Cumulativv Reward
4000
2,000
2000
500
1000
1500
2000
1000
1500
500
1000
1500
2000
Steps (x10P))
Steps (x10F)
Steps (x(10))
HalfCheetahPOMDDPv2
HopperPOMDP.v3
Cumulative Reward
Reward
8000
10000
2000
1000
Steps (x10')
Steps (x10')
TD3 SPTT-S
TD3-RS (Ouns
SAC-BFTTS
TD3-vanila
TD3-oracle
SAC-Windows
TD3: WindowS5 TD3-Wrrdow-32
T03-8PTT-30
SAC-RS (Oun) SAC-BPT-322
SACvanill
SACoraace
SAC-Window-2
Figure 12: SAC and results on the MuloCo Gym Tasks with partially observable policies and privileged critics. Additionally, masses of each link are randomly sampled at the beginning of each episode. The number behind the BPTT implementations indicates the truncation length. Abscissa shows the cumulative reward. Ordinate shows the training time in hours.
Privileged Information
No Privileged Information
35-
0.8-
Normalized Return
00.8
mime Time (5)
3.0
0.6-
0.6
20-
10-
0.2
0.2
TD3
SLAC
0.8
25-
Mormalized Return
0.88
Traiine Time [h]
0.6
0.6
155
Training
0.4
0.4
0.2-
0.0
SAC
TD3
PPO
SLAC
PPO
RS (Ours)
BPTT-5
BPTT-32
Window-5
Window-32
Vanilla
Oracle
Figure 13: Results show the mean and 95% confidence interval of the normalized return and the training time needed for million steps across high-dimensional TOP row; Ant and Humanoid and low-dimensional
BOTTOM row: Hopper. Walker and HalfCheetah POMDP Gym tasks. Comparison of different RL agents on POMDP tasks using our stateful gradient estimator, BPTT with truncation length of 32, SLAC algorithm and stateless versions of the algorithms using a window of observation with length 5 and 32. The "Oracle" approach is a vanilla version of the algorithm using full-state information.
35

---

Published as a conference paper at ICLR 2024
Table I: Runtime comparison of different algorithms on all reinforcement learning tasks. For SAC and TD3, privileged information for the critic is used, and the tasks include randomization of the the velocities, while for PPO no privileged information is used and only occlusion of the velocities used for the tasks. The window, vamilla, and oracle approaches are grayed out as they do not constitute stateful policies. The bold values indicate the lowest training time across all stateful algorithms. All values are computed hy taking the mean over 10 seeds and computing the 95% confidence interval.
and occlusion of
Time in hours needed for I million environment steps [lower is better]
BPTT
Window
Vanilla
Oraele
(Ours)
Env.
5
32
32
Stateful Policy
yes
yes
n10
n1
no
n1
yes
7.8
4.8 0.120
10.6
2.5 0.092
0.256
Hopper
6.8 = 0.168
22.3 0,406
3.5 0.036
4.08:.008
2.3 0.026
2.2 0.070
5.9 : 0.017
SAC
HallChectah
6.9 0.291
27.1 + 0.361
0.054
5.14-0.106
2.3 0.071
2,4 2.0000000
6.0 + 0.041
Humanoid
10.1 0.285
32.9 1.317
0.118
0
3.3 3.049
7,9 + 0.027
Walker2d
6,8 + 0,202
22.7 + 0.696
3.7 + 0.051
0.104
2.4 0.034
2,3 + 0.067
6.0 + 0.026
Ant
=: 0.019
0.119
0.017
0.127
0.018
0.010
3.0 3.0 0.124
Hopper
2.6 0.058
4.5 0.176
1.9 0.054
0.030
0.025
1.0:1.0018
2.8 - 0.093
TD3
HalfCheetah
0.010
5.8
2.0 0.086
0.070
1.3=0.002
122+:.044
2.7 + 0.100
Humanoid
0.04
0.325
2.7 0.047
0.16
18
1.8
3.6 0.053
Walker2d
2.4 0.006
5.1 0.242
1,9 0.018
2.2 0.062
1,3
1.3 0.012
2.9 + 0.052
Ant
2,8 + 0.001
3.3 0.001
0.7 + 0.088
1.079.079
0.4 0.005
0.6 + 0.007
2.1 + 0.036
Hopper
2.5
3.1 0.001
0.5 = 0.004
0,5
0.3 0.006
0.4 1 0.004
1.8 : 0.061
PPO
HallChestah
2.5 0.001
3.1 0.001
0.5 + 0.003
0.6 0.000
0.4 0.012
0.4 + 0.001
2.2 0.044
Humanoid
2.9 0.000
3.5 0.001
3.5 = 0.001
3,5 0.001
0,5 1: 0.052
0.7 1 0.009
2.1 + 0.033
Walker2d
2.5
0.5
0.041
Table 2: Runtime comparison of different algorithms on all reinforcement learning tasks. For SAC and TD3, privileged information for the critic is used, and the tasks include randomization of the masses and occlusion of the velocities, while for PPO no privileged information is used and only occlusion of the velocities is used for the tasks. The window, vamilla, and oracle approaches are grayed out as they do not constitute stateful policies. The factors are computed for million environment steps. The bold values indicate the lowest factor across all stateful algorithms. All values are computed by taking the mean over 10 seeds.
Factor time increase compared to Vanilla version [lower better]
BPTT
Window
Vanilla
RS (Ours)
Env.
32
32
Stateful Policy
yes
yes
yes
Ant
2.804
10.935
1.730
3.822
1.000
2.241
Hopper
2.935
9.599
1.497
1.956
1.000
2.531
HalfCheetah
3.022
1.614
2.213
1.000
Humanoid
2.936
9.580
1.887
4.106
1.000
2.290
Walker2d
2,878
9,597
1.578
1,948
1,000
2,526
Ant
2.024
4.535
1,546
2.620
1.000
2.056
Hopper
2.012
3,476
1.489
1.643
1.000
2.146
TD3
HallCheetah
2.069
4.606
1.555
1.864
1.000
2.120
Humanoid
2.037
4.555
1,512
2.896
1,000
2.024
Walker2d
1.913
4.024
1.527
1.724
1.000
2.245
Ant
7.845
9.457
2.120
4.637
1.000
5.815
Hopper
9.12%
11.194
1.706
1.825
1.000
6.602
HalfCheetah
6.888
8.454
1.331
1.536
1.000
6.023
Humanoid
6.018
7.183
7.183
7.183
1.000
4.264
Walker2d
8.735
10.698
1.666
1.829
1.000
6.048
36

---

Published as a conference paper at ICLR 2024
AtlasPOMDP-Carry
HumanoidPOMDP-Walk
HumanoidPOMDP-Run
Cumulative Rewara
1000
Cumulative Reward
800
750
750
600
500
500
Cumulative
400
250
200
5000
10000
15000
5000
10000
15000
5000
10000
15000
Steps (x101)
Steps (x1D*)
Steps (x10*)
GAIL-PTT-S
GAIL-6PT-32
GAILARS (Ours)
GAIL-vanilla
Figure 14: GAIL results all IL tasks with partially observable policies and privileged critics. The number behind the BPTT implementations indicates the truncation length. Abscissa shows the cumulative reward. Ordinate shows the number of training steps (x10").
AtlasPoMDP-Carry
HumanoidPOMDP-Wal.
HumanoidPOMDP-Rum
Cumulative Reward
Cumulative Reward
750
600
40.0
500
400
250
200
500
1000
1500
2000
1000
1500
2000
500
1000
1500
2000
Steps (x10²)
Steps (x10²)
Steps (x107)
LSIQ-8PTT-
LSIQ-6PTT-32
LSIO-R5 (Ours)
LSIO-vanilla
Figure 15: LS-IQ results all IL tasks with partially observable policies and privileged crities. The number behind the BPTT implementations indicates the truncation length. Abscissa shows the cumulative reward. Ordinate shows the number of training steps (x10").
E.5 IMITATION LEARNING EXPERIMENTS
In this section, we present further results of the tasks from the main paper. As can seen in Figure [4] and 15 our results show that the S2PG version of GAIL and LS-IQ GAIL-RS and LSIQ-RS are competitive with the BPTT variants, even when the truncation length is increased. In the HumanoidPOMDP-WWll task, GAIL-RS is even able to outperform by a large margin the BPTT variant. As stressed in the paper, the advantage of using S2PG over BPTT is the reduced computation time. As LS-IQ runs SAC in the inner RL loop, and GAIL runs PPO, the runtimes of these algorithms are very similar to their RL counterparts. For the Atlas experiments, it can be seen that the vanilla policy GAIL already performing well, despite being outperformed by our method. This shows that this task does not require much memory. In contrast, the humanoid tasks require much more memory, as can be seen by the poor results of the vanilla approaches. Interestingly, found the GAIL-BPTT variants often very difficult to train these tasks, which is the reason why they often result in unstable learning.
37


Reviews:
[{"review_id": "CLfiXkdLA7", "reviewer": "Reviewer_gHcm", "summary": "The paper addresses the problem of learning policies with long term  history. Traditional methods often employ recurrent architectures, which rely on a persistent latent state that is modified at every time increment. However, this presents a challenge: when calculating the gradient of a loss function in such architectures, the gradient needs to be back-propagated through all preceding time steps. This process can lead to either vanishing or exploding gradients, making the training of recurrent models particularly difficult, especially as the historical data size increases. To address this issue, the authors introduce an alternative approach wherein the model's internal state is represented as a stochastic variable that is sampled at each time step. As a result, the state's stochastic nature prevents the direct computation of an analytical gradient, thereby circumventing the issues associated with backpropagation over time. The paper goes on to adapt established theoretical frameworks to this new model and suggests a method for incorporating actor-critic techniques. Empirical validation is conducted on a range of environments that are structured as Partially Observable Markov Decision Processes (POMDPs) by omitting certain observations.  It is shown that the proposed model achieves reasonable performance in comparison to BPTT-based approaches.", "questions": "(see previous comments)", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well-written and easy to follow. The supplementary appendix, which was not examined in detail, appears to be a valuable extension of the main text. The concept of characterizing the policy's internal state as a stochastic variable is intriguing and yields a sophisticated formulation. Additionally, the paper offers robust theoretical contributions and presents a methodology to modify conventional algorithms to encompass this concept.", "weaknesses": "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. \n\nThen, 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\n\nUsing stochastic variables as state of a policy is something made typically when using Thompson sampling-like methods. Papers like \"Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables\" are also using a stochastic internal state. How do you position your work w.r.t these approaches ? \n\nIn the experiments, it is not clear how the z distribution is modeled, and there is no discussion about possible choices and their possible impact. For instance, what about using a multinomial distribution? Discussing that point would be interesting.\n\nFigure 1 is misleading since there are no arrows between the s nodes and the z nodes in the graph on the right and it seems that the sequence of z does not depend on the observations", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kSU4SnnzT7", "reviewer": "Reviewer_HaqJ", "summary": "The paper proposes a novel method for decomposing a stateful policy into a stochastic internal state kernel and a stateless policy, resulting in a new policy gradient that is applicable to POMDPs without the need for backpropagating through time (BPTT). At the heart of this technique is the modification of the policy to output not only an action but also a prediction over the subsequent internal state. The authors have derived both stochastic and deterministic policy gradient theorems for this framework and have expanded the variance analysis of Papini et al. for policy gradient estimators. The experimental results demonstrate that the proposed method rivals algorithms that use full BPTT while requiring considerably less computational effort.", "questions": "- The occupancy measure is introduced without defining $z$\n\n- Is the initial internal state learned, or is it initialized to zero at the beginning of each episode?\n\n- 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?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper is very clear and provides a thorough presentation of the theory proposed by the method. I believe that this work will serve as a good reference for any work building on alternatives to BPTT for POMDPs.\n\n- To my knowledge, the results presented are novel.\n\n- The authors have conducted a detailed analysis of the algorithm across several complex tasks.\n\n- While the theory occasionally presents straightforward extensions of classical policy gradient results, it is explained with exceptional clarity both in the main text and in the appendix.", "weaknesses": "- 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.\n\n- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TGayVuXkqF", "reviewer": "Reviewer_wQTz", "summary": "The paper presents a new method entitled 'stochastic stateful policies' for performing reinforcement learning (RL) with policies having an internal state. The method can be applied to POMDP settings and other non-Markovian setups where it is necessary to have access to the whole trajectory, not only the current state. The method is compared against the state-of-the-art backpropagation through time (BPTT) approach and is more computationally efficient but may result in higher variance. Fundamental theorems are provided, and the method is evaluated in a series of problems, showing improvement, especially in higher dimensional cases.", "questions": "### Main questions\n* Provide a wall-time comparison of the introduced S2PG approach with (as I understand more time-consuming) BPTT approach.\n* Used assumptions in the theoretical part - argue their physical motivation and preferably cite some established works that introduced them,\n* Hiding velocity experiment - this will not check the benefits of any long-term history encoding, as only two states can be used to approximate velocity.\n* I have a suggestion for another experiment - simplify further the policy by removing the $\\pi_\\theta^a$ , and keeping only $\\pi_\\theta^z$. \n* What do you mean by  \"our method could be used in combination with complex policy structures such as Neural ODEs.\" can you elaborate?\n\n### Minor remarks/questions\n* p. 3 eq. (1) provide formula for $J(\\tau)$;\n* p. 4, caption Fig. 1 Do you rather mean 'from left to right' ?\n* p. 4 'equations equation' -> 'equation;\n* p. 6 \"causality dilemma: On one hand, the policy can learn a condensed representation of the\nhistory assuming accurate Q-value estimates. On the other hand, the Q-function can be learned using bootstrapping techniques, assuming a reliable condensed representation of the history. But doing\nboth at the same time often results in unstable training.\" be more clear here;\n* p. 6 \"We extend the letter\" -> the latter;\n* p. 8 caption Fig. 2 \"for 1 Mio. steps\" -> what is Mio. ?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "Designing new RL approaches tailored for the POMDP setting is important, as a practical deployment will certainly be non-Markovian. To my knowledge, the existing state-of-the-art is to use existing RL algorithm with stateful actors/critics, e.g., along the current state, a history encoded using an LSTM/RNN/GRU is provided.\n\nThe paper has much content; the theoretical section is rich, with a few fundamental results -- the policy gradient theorem, the deterministic policy gradient theorem, and IMHO, the most important - a variance upper bounds are provided in the studied setting. However, it mainly adapts existing results to the current setting, which is unsurprising. Clean proofs of the theorems are provided. \nThe experimental part is also rich; a few experiments reported (unfortunately mostly in the appendix) demonstrating the approach's potential. I think it is fair to say that the paper contains a thorough experimental study.\n\nThe method performs better with less computational and memory burden than the SOTA BPTT approach in high-dimensional problems. Moreover, the authors are honest, and some settings in which the method does not perform well (likely due to increased variance) are also provided, which is nice. \n\nOverall, after reading the main part of the paper, I saw the paper as a borderline paper. However, luckily, after going through the much longer appendix, I became more positive about the paper. Let me note here that as a reviewer, I am not obliged to make a detailed pass over the paper appendix and should base my judgment upon the main part content. \n\nConsidering the whole content, this is a convincing paper. However, there is a danger if a reader goes through the main part only may not appreciate the results fully. This brings my main concern about the paper - I am not sure if the authors took significant effort in making the main part of the paper stand-alone and convincing enough, considering the amount and quality of produced results. The main part needs a major revision before publication. I will reconsider my score after my concerns are addressed in the rebuttal phase. See detailed remarks and questions below.", "weaknesses": "My main concern with the paper is that the main part needs to present the available results self-contained and convincing enough, which is not the case. Before accepting the paper, it will need thorough edits addressing the concerns and the questions I present below.\n\n1. 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.\n\n2. 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\".\n\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)