Policy Gradient With Serial Markov Chain Reasoning
Abstract
We introduce a new framework that performs decision-making in reinforcement learning (RL) as an iterative reasoning process.
We model agent behavior as the steady-state distribution of a parameterized reasoning Markov chain (RMC), optimized with a new tractable estimate of the policy gradient.
We perform action selection by simulating the RMC for enough reasoning steps to approach its steadystate distribution.
We show our framework has several useful properties that are inherently missing from traditional RL.
For instance, it allows agent behavior to approximate any continuous distribution over actions by parameterizing the RMC with a simple Gaussian transition function.
Moreover, the number of reasoning steps to reach convergence can scale adaptively with the difficulty of each action selection decision and can be accelerated by re-using past solutions.
Our resulting algorithm achieves state-of-the-art performance in popular Mujoco and DeepMind
Control benchmarks, both for proprioceptive and pixel-based tasks.
Introduction
Reinforcement learning (RL) has the potential to provide a general and effective solution to many modern challenges.
Recently, this class of methods achieved numerous impressive milestones in different problem domains, such as games [1–3], robotics [4–6] and other meaningful real-world applications [7–9].
However, all these achievements relied on massive amounts of data, controlled environments, and domain-specific tuning.
These commonalities highlight some of the current practical limitations that prevent RL to be widely applicable [10].
In the deep RL framework, practitioners train agents with the end goal of obtaining optimal behavior .
Traditionally, agent behavior is modeled with feed-forward policies regressing from any state to a corresponding distribution over actions.
Such formulation yields practical training objectives in both off-policy [11–13] and on-policy settings [14–16].
However, we identify three inherent properties of this rigid representation of behavior that could considerably impact expressivity and efficiency in continuous control tasks.
First, agent behavior is restricted to a class of tractable distributions, which might fail to capture the necessary complexity and multi-modality of a task.
Second, the policy performs a fixed reasoning process with a feed-forward computation, which potency cannot adapt to the varying complexity of individual action selection problems.
Third, decision-making is performed every time from scratch , without re-using any past information that might still inform and facilitate the current action selection problem.
Unlike RL policies, human reasoning does not appear to follow a rigid feed-forward structure.
In fact, a range of popular psychological models characterize human decision-making as a sequential process with adaptive temporal dynamics [17–20].
Many of these models have found empirical groundings in neuroscience [21–24] and have shown to effectively complement RL for capturing human behavior in experimental settings [25, 26].
Partly inspired by these works, we attempt to reframe the deep RL framework by making use of a similar flexible model of agent behavior, in order to counteract its aforementioned limitations.
We introduce serial Markov chain reasoning - a new powerful framework for representing agent behavior.
Our framework treats decision-making as an adaptive reasoning process , where the agent sequentially updates its beliefs regarding which action to execute in a series of reasoning steps .
We model this process by replacing the traditional policy with a parameterized transition function , which defines a reasoning Markov chain (RMC).
The steady-state distribution of the RMC represents the distribution of agent behavior after performing enough reasoning for decision-making.
Our framework naturally overcomes the aforementioned limitations of traditional RL.
In particular, we show that our agent’s behavior can approximate any arbitrary distribution even with simple parameterized transition functions.
Moreover, the required number of reasoning steps adaptively scales with the difficulty of individual action selection problems and can be accelerated by re-using samples from similar RMCs.
To optimize behavior modeled by the steady-state distribution of the RMC, we derive a new tractable method to estimate the policy gradient .
Hence, we implement a new effective off-policy algorithm for maximum entropy reinforcement learning (MaxEnt RL) [27, 28], named Steady-State Policy
Gradient (SSPG).
Using SSPG, we empirically validate the conceptual properties of our framework over traditional MaxEnt RL.
Moreover, we obtain state-of-the-art results for popular benchmarks from the OpenAI Gym Mujoco suite [29] and the DeepMind Control suite from pixels [30].
In summary, this work makes the following key contributions:
We propose serial Markov Chain reasoning a framework to represent agent behavior that can overcome expressivity and efficiency limitations inherent to traditional reinforcement learning.
Based on our framework, we derive SSPG , a new tractable off-policy algorithm for MaxEnt RL.
We provide experimental results validating theorized properties of serial Markov Chain reasoning and displaying state-of-the-art performance on the Mujoco and DeepMind Control suites.
Background 2.1 Reinforcement learning problem
We consider the classical formulation of the reinforcement learning (RL) problem setting as a Markov
Decision Process (MDP) [31], defined by the tuple ( S, A, P, p 0 , r, γ ) .
In particular, at each discrete time step t the agent experiences a state from the environment’s state-space, s t ∈ S , based on which it selects an action from its own action space, a t ∈ A .
In continuous control problems (considered in this work), the action space is typically a compact subset of an Euclidean space R dim ( A ) .
The evolution of the environment’s state through time is determined by the transition dynamics and initial state distribution, P and p 0 .
Lastly, the reward function r represents the immediate level of progress for any state-action tuple towards solving a target task.
The agent’s behavior is represented by a state-conditioned parameterized policy distribution π θ .
Hence, its interaction with the environment produces trajectories, τ = ( s 0 , a 0 , s 1 , ..., s T , a T ) , according to a factored joint distribution p π θ ( τ ) =
p 0 ( s 0 ) (cid:81) Tt =
π θ ( a t | s t ) P ( s t +1 | s t , a t ) .
The RL objective is to optimize agent behavior as to maximize the discounted sum of expected future rewards: arg max θ E p
πθ ( τ ) (cid:104)(cid:80) Tt =
γ t r ( s t , a t ) (cid:105) .
2.2 Maximum entropy reinforcement learning and inference
Maximum entropy reinforcement learning (MaxEnt RL) [32] considers optimizing agent behavior for a different objective that naturally arises when formulating action selection as an inference problem [33–36].
Following Levine [28], we consider modeling a set of binary optimality random variables with realization probability proportional to the exponentiated rewards scaled by the temperature α , p ( O t | s t , a t ) ∝ exp( 1 α r ( s t , a t )) .
The goal of MaxEnt RL is to minimize the KL-divergence between trajectories stemming from agent behavior, p π θ ( τ ) , and the inferred optimal behavior, p ( τ | O 0: T ) :
[Equation]
The resulting entropy-regularized objective introduces an explicit trade-off between exploitation and exploration, regulated by the temperature parameter
α scaling the policy’s entropy.
An effective choice to optimize this objective is to learn an auxiliary parameterized soft Q-function [37]:
[Equation]
Given some state, Q πϕ ( s, · ) represents an energy-function based on the expected immediate reward and the agent’s future likelihood of optimality from performing any action.
Thus, we can locally optimize the MaxEnt objective by reducing the KL-divergence between π and the canonical distribution of its current soft Q-function.
This is equivalent to maximizing the expected soft Q-function’s value corrected by the policy’s entropy, resembling a regularized policy gradient objective [11, 12]:
[Equation]
The policy is usually modeled with a neural network outputting the parameters of some tractable distribution, such as a factorized Gaussian, π θ ( ·| s ) =
N ( µ θ ( s ); Σ θ ( s )) .
This practice allows to efficiently approximate the gradients from Eqn. 3 via the reparameterization trick [38].
We consider the off-policy RL setting, where the agent alternates learning with storing new experience in a data buffer, D .
We refer the reader to Haarnoja et al.
[13, 39] for further derivation and practical details.
Policy Gradient with serial reasoning 3.1 Reasoning as a Markov chain
We introduce Serial Markov Chain Reasoning , a new framework to model agent behavior, based on conceptualizing action selection as an adaptive, sequential process which we refer to as reasoning .
Instead of using a traditional policy, the agent selects which action to execute by maintaining an internal action-belief and a belief transition (BT-) policy , π b ( a ′ | a, s ) .
During the reasoning process , the agent updates its action-belief for a series of reasoning steps by sampling a new action with the BT-policy π b taking both environment state and previous action-belief as input.
We naturally represent this process with a reasoning Markov chain (RMC), a discrete-time Markov chain over different action-beliefs, with transition dynamics given by the BT-policy.
Hence, for any input environment state s and initial action-belief a 0 , the n-step transition probabilities of the RMC for future reasoning steps n = 1 , 2 , 3 , ... are defined as:
[Equation]
for n >
[Equation] [Equation]
Given a compact action space and a BT-policy with a non-zero infimum density, we can ensure that as the number of reasoning steps grows, the probability of any action-belief in the RMC converges to some steady-state probability which is independent of the initial action-belief.
1 We denote this implicit probability distribution as the steady-state (SS-) policy , symbolized by π s ( a | s ) :
Lemma 3.1.
Steady-state convergence.
For any environment state s , consider a reasoning Markov chain (RMC) defined on a compact action space A with transition probabilities given by π b ( a ′ | a, s ) .
Suppose that inf { π b ( a ′ | a, s ) : a ′ , a ∈ A } > 0 .
Then there exists a steady-state probability distribution function π s ( ·| s ) such that:
[Equation]
Proof.
See Appendix A.
The RMC’s steady-state probabilities can be interpreted as representing the distribution of agent’s behavior after an appropriate number of reasoning steps are performed.
In this work, we strive to optimize the agent’s behavior following the MaxEnt RL framework described in Section 2.
In particular, we consider learning a parameterized BT-policy, π bθ , to produce appropriate transition probabilities for each environment state such that the SS-policy, π sθ , from the resulting RMC optimizes:
[Equation]
Here, Q sϕ is a parameterized soft Q-function for the agent’s behavior from π s , which we learn by minimizing a squared soft Bellman loss utilizing delayed parameters ϕ ′ and samples from π sθ : arg min ϕ
[Equation]
In Fig.
2, we illustrate the relationship between a learned BT-policy, the corresponding SS-policy, and the soft Q-function in a 1-dimensional toy task (see App.
C for details).
In this example, the
BT-policy is parameterized as a simple squashed Gaussian distribution, with unimodal transitions between consecutive action beliefs (Fig. 2, Left).
We obtain samples of agent behavior (the SS-policy) by performing a series of reasoning steps, using the BT-policy to simulate the RMC until we approach steady-state convergence.
By plotting the resulting empirical distribution of agent behavior, we see it closely matches the multi-modal, non-Gaussian canonical distribution from its soft Q-function (Fig. 2, Right).
This example shows how the expressive power of agent behavior in our framework can go far beyond the BT-policy’s simple parameterization, enabling for the effective maximization of complex and multi-modal MaxEnt objectives.
3.2 Learning the belief transition policy
We propose a new method to estimate the policy gradient of the BT-policy, π bθ , for optimizing the steady-state MaxEnt objective described in Section 3.1.
We note that the gradient from Eq.
6 involves differentiating through an expectation of the steady-state policy, π sθ .
However, π sθ is only implicitly defined, and its connection with the actual BT-policy or its parameters does not have a tractable closed-form expression.
To approach this problem, we introduce a family of n-step extensions to the soft Q-function, Q sn :
S × A (cid:55)→ R for n = 0 , 1 , 2 , . . . , defined as:
[Equation]
with ∇ θ Q sn ( s, a ) = 0 .
Intuitively, each n-step soft Q-function Q sn ( s, a ) outputs the expected soft Q-value after performing n reasoning steps in the RMC from the initial action-belief a .
However, we treat the output of each n-step soft Q-function as being independent of the actual parameters of the BT-policy, θ .
Hence, we can interpret computing Q sn ( s, a ) as simulating the RMC with a fixed and immutable copy of the current π bθ .
We use this definition to provide a convenient notation in the following new Theorem that expresses the policy gradient without differentiating through π sθ :
Theorem 3.2.
Steady-state policy gradient.
Let π bθ ( ·| a, s ) be a parameterized belief transition policy which defines a reasoning Markov chain with a stationary distribution given by the steady-state policy π sθ ( ·| s ) .
Let Q s be a real function defined on S × A , with a family of n-step extensions { Q sn } as defined in Eq. 8.
Suppose π b , Q s and their gradient with respect to θ (denoted ∇ θ ) are continuous and bounded functions.
Then
[Equation]
Proof.
See Appendix A.
Using Lemma 3.1 ( steady-state convergence ), we can approximate the policy gradient expression in Eq.
9 with an arbitrarily small expected error using a finite number of n-step soft Q-functions, i.e., N (see App.
A).
An intuition for this property follows from the fact that for large enough n , Lemma 3.1 implies that π bn ( a | a ′ , s )
≈ π sθ ( a | s ) and, thus, Q sn ( s, a ′ )
≈ (cid:82)
A π bθ s ( a | s ) Q sϕ ( s, a ) da .
Therefore, the value of each Q sn ( s, a ′ ) will be independent of the BT-policy’s action a ′ , such that ∇ θ E a ′ ∼ π bθ ( ·| a,s )
[ Q s n ( s, a ′ )] ≈ 0 .
In other words, each subsequent step in the RMC introduces additional randomness that is independent of a ′ , causing a warranted vanishing gradient phenomenon [41] which culminates with converging to π sθ .
Using a similar notation as Haarnoja et al.
[39], we apply the reparameterization trick [38] to express the BT-policy in terms of a deterministic function f bθ ( a, s, ϵ ) , taking as input a Gaussian noise vector ϵ .
This allows to rewrite the gradient in each inner expectation in the sum from Eq.
9 as:
[Equation]
where a 0 = f bθ ( a, s, ϵ 0 ) .
We can apply the same reparameterization for all n-step soft Q-functions, to establish a new relationship between the gradient terms ∇ a 0 Q sn ( s, a 0 ) :
[Equation]
where, a 1 = f ( a 0 , s, ϵ 1 ) .
In Eq.
11, we purposefully omit the dependence of f b and π b from θ since each Q sn term is a local approximation of the RMC that does not depend on θ (as defined in Eq.
By recursively applying this relationship (Eq. 11) to ∇ a 1 Q sn − 1 ( s, a 1 ) and all subsequent gradient terms we obtain:
[Equation]
12) where a i = f b ( a i − 1 , s, ϵ i ) for i = 1 , . . . , n .
By combining Eq.
10 and Eq.
12, we can thus reparameterize and express the whole sum in Eq.
9 as:
[Equation]
Eq.
13 intuitively corresponds to differentiating through each Q s n ( s, a ′ ) term by reparameterizing the
RMC .
Hence, to get a sample estimate of the policy gradient we can simulate the reparameterized
RMC for N reasoning steps to obtain a 1 , ..., a N , compute each Q sϕ ( s, a n ) term, and backpropagate (e.g., with autodifferentiation ).
Following Haarnoja et al.
[13, 39], we can apply Theorem 3.2 and easily extend the same methodology to estimate the MaxEnt policy gradient from Eq. 6 that also involves an extra entropy term.
We include this alternative derivation in App.
A for completeness.
input: s , current state a 0 ∼ ˆ A
N ← 0 R p ← + ∞ while R p > 1 .
1 do a N + 1 ∼ π bθ ( ·| a N ) N ← N +
Update R p with a 1 : N ▷ Eq.16 ˆ N ←
ρ ˆ N + (1 − ρ ) N ▷ ρ ∈ [0 , 1) ˆ
A ← ˆ A ∪ a 1 : N output: a ∼ a 1 : N
Algorithm input: D , data buffer ( s, a, s ′ , r ) ∼ D a 0 ∼ π bθ ( ·| a, s ′ ) for n ← 0 , ⌈ ˆ N ⌉ do
Q sn ← Q sϕ ( s ′ , a n )
▷ Eq.ϵ n +1 ∼ N (0 , 1) ,
a n +1 = f b ( a n , s, ϵ n +1 ) ∇ θ Q sϕ ← ∇ θ ( (cid:80) ⌈ N ⌉ n =0 Q sn )
▷ Thm.
3.2 arg min θ J ( θ ) ▷ Eq.
6 a ′ ∼ a 1: ⌈ ˆ N ⌉ arg min ϕ J ( ϕ )
▷ Eq.
3.3 Action selection and temporal consistency
To collect experience in the environment, we propose to perform reasoning with the BT-policy starting from a set of different initial action-beliefs { a 0 0 , ..., a M 0 } .
We batch this set as a single input matrix, a 0 , to make effective use of parallel computation.
To reduce the required number of reasoning steps and facilitate detecting convergence to π sθ , we identify two desirable properties for the distribution of action-beliefs in a 0 .
In particular, initial action-beliefs should 1) be likely under π sθ , and 2) cover diverse modes of π sθ .
Property (1) should logically accelerate reasoning by providing the BT-policy with already-useful information about optimal behavior.
Property (2) serves to provide the BT-policy with initial information of diverse behavior, which facilitates convergence detection (Sec. 3.4) and expedites reasoning even if the RMC has slow mixing times between multiple modes.
To satisfy these properties, we use a simple effective heuristic based on common temporal-consistency properties of
MDPs [42, 43].
Especially in continuous environments, actions tend to have small individual effects, making them likely relevant also for environment states experienced in the near future.
Thus, we propose storing past action-beliefs in a fixed sized buffer, called the short-term action memory , ˆ A , and use them to construct a 0 .
We find this strategy allows to effectively regulate the initial action-beliefs quality and diversity through the size of ˆ A , accelerating convergence at negligible cost.
3.4 Detecting convergence to the steady-state policy
A key requirement for learning and acting with BT-policies, as described in Sections 3.2 and 3.3, is the ability to determine a sufficient number of reasoning steps ( N ) for the action-belief distribution to converge.
Given the properties of the RMC, there exist different analytical methods that provide a priori bounds on the rate of convergence [44–46].
However, using any fixed N would be extremely limiting as we expect the BT-policy and the properties of its resulting RMCs to continuously evolve during training.
Moreover, different tasks, states, and initial action-beliefs might affect the number of reasoning steps required for convergence due to different levels of complexity for the relative decisionmaking problems.
To account for similar conditions, in the Markov Chain Monte Carlo literature, the predominant approach is to perform a statistical analysis of the properties of the simulated chain, choosing from several established convergence diagnostic tools [47–49].
Hence, we propose to employ a similar adaptive strategy by analyzing the history of the simulated RMC to determine the appropriate number of reasoning steps.
Since we apply π bθ from a diverse set of initial action beliefs (see Section 3.3), we base our convergence-detection strategy on the seminal Gelman-Rubin (GR) diagnostic [50] and its multivariate extension [51].
In particular, the multivariate GR diagnostic computes the pseudo scale reduction factor (PSRF), a score representing whether the statistics of a multivariate variable of interest have converged to the steady-state distribution.
The intuition behind this diagnostic is to compare two different estimators of the covariance for the unknown steady-state distribution, making use of either the samples within each individual chain and between all different chains.
Thus, as the individual chains approach the true steady-state distribution, the two estimates should expectedly get closer to each other.
The PSRF measures this precise similarity based on the largest eigenvalue of their matrix product.
For our use-case, we employ the PSRF to determine the convergence of the set of action-beliefs a 1 : N , as we perform consecutive reasoning steps with π bθ .
Following [51], we calculate the average sample covariance of the action-beliefs within each of the parallel chains ( W ) computed from a batched set of initial action-beliefs a
=
[ a 10 , a 20 , . . .
a M 0 ] :
[Equation]
We compare W with an unbiased estimate of the target covariance, constructed from the sample covariance between the different parallel chains ( B ):
[Equation]
The PSRF for a 1 : N is then computed from the largest eigenvalue ( λ max ) of the product W − 1 B , as:
[Equation]
Thus, as the individual chains approach the distribution of π sθ , the PSRF ( R p ) will approach 1.
Following Brooks and Gelman [51], we use R p < 1 . 1 as an effective criterion for determining the convergence of a 1 : N .
In practice, we also keep a running mean of the current number of reasoning steps for convergence, ˆ N .
We use ⌈ ˆ N ⌉ as the number of reasoning steps to simulate the RMC with π bθ when computing gradients from Eqs.
⌈ ˆ N ⌉ is a safe choice to ensure near unbiased optimization since R p < 1 .
1 is considered a very conservative criterion [52] and we can learn by simulating the RMC from recent actions stored in the data buffer, which are already likely close to optimal.
We provide further details regarding our implementation and its rationale in App.
B. We provide a simplified summary of our adaptive reasoning process for acting and learning in Algs.
3.5 Advantages of serial Markov chain reasoning
Based on the above specification, we identify three main conceptual advantages of our serial Markov chain reasoning framework.
Unlimited expressiveness.
The distribution of agent behavior given by the SS-policy π sθ , is a mixture model with potentially infinitely many components.
Thus, even a simple Gaussian parameterization of the BT-policy π bθ would make π sθ a universal approximator of densities , providing unlimited expressive power to the agent [53, 54].
Adaptive computation.
The number of reasoning steps performed to reach approximate convergence is determined by the properties of each environment state’s RMC.
Hence, the agent can flexibly spend different amounts of computation time based on the complexity of each action-selection problem, with potential gains in both precision and efficiency.
Information reuse.
By storing past solutions to similar RMCs, we can initialize the reasoning process with initial action-beliefs that are already close to π sθ .
This allows using the temporal-consistency properties of the MDP to exploit traditionally discarded information and accelerate agent reasoning.
We provide empirical validation for these properties in Section 4.2.
Experimentation 4.1 Performance evaluation
We evaluate the serial Markov chain reasoning framework by comparing its performance with current state-of-the-art baselines based on traditional RL.
We consider 6 challenging Mujoco tasks from
Gym [29, 56] and 12 tasks pixel-based tasks from the DeepMind Control Suite (DMC) [30].
In both settings, we base our implementation on MaxEnt RL, replacing the traditional policy with a Gaussian
BT-policy optimized with the training procedures specified in Sec. 3.
Other orthogonal design choices (e.g., network architectures) follow contemporary RL practices, we refer to App.
C or the code for full details.
We call the resulting algorithm Steady-State Policy Gradient (SSPG).
We report the mean performance curves and aggregate metrics using the statistical tools from Rliable [55].
In particular, we compare normalized performance profiles [57], interquantile mean (IQM), and probability of improvements over baselines with the Mann-Whitney U statistic [58].
The reported ranges/shaded regions represent 95% stratified bootstrap confidence intervals (CIs)
In App.
D, we provide per-task results and further statistical analysis.
For each experiment, we collect the returns of SSPG over five seeds, by performing 100 evaluation rollouts during the last 5% of steps.
Mujoco suite.
We evaluate on a challenging set of Mujoco tasks popular in recent literature.
We compare SSPG with recent RL algorithms achieving state-of-the-art sample-efficiency performance on these tasks, which utilize large critic ensembles and high update-to-data (UTD) ratios.
We consider
REDQ [60] and MBPO [61] for state-of-the-art algorithms based on the traditional model-free and model-based RL frameworks.
We also compare with iterative amortized policy optimization (IAPO)
[62], in which the agent performs iterative amortization to optimize its policy distribution [63].
This procedure for action selection is more computationally involved than our agent’s reasoning process , as it requires both evaluating the policy and computing gradients at several iterations.
Yet, as IAPO is still based on the traditional policy gradient framework, its benefits are solely due to reducing the amortization gap with an alternative action inference procedure.
To ground different results, we also show the performance of the seminal Soft Actor-Critic (SAC) algorithm [39], upon which all considered policy gradient baselines are based on.
To account for the additional computational cost of training an agent with serial Markov chain reasoning, we use a UTD ratio that is half the other algorithms.
On our hardware, this makes SSPG faster than all other modern baselines (see App. D).
Figure 3 (Top) shows the performance results after 100K environment steps.
Individual scores are normalized using the performance of SAC after 3M steps, enough to reach convergence in most tasks.
SSPG considerably outperforms all prior algorithms with statistically meaningful gains, as per the conservative Neyman-Pearson statistical testing criterion [64].
Furthermore, SSPG even stochastically dominates all considered state-of-the-art baselines [65].
We obtain similar results evaluating at 50K and 200K steps (App. D).
In comparison, IAPO obtains lower performance than other non-iterative baselines while being the most compute-intensive algorithm.
This indicates that, for sample-efficiency, only reducing the amortization gap beyond direct estimation might not provide significant benefits.
Instead, serial Markov chain reasoning’s improved expressivity and flexibility appear to considerably accelerate learning, yielding state-of-the-art performance in complex tasks.
DeepMind Control suite.
To validate the generality of our framework, we also evaluate on a considerably different set of problems: 12 pixel-based DMC tasks.
We follow the recent task specifications and evaluation protocols introduced by Yarats et al.
We compare SSPG with
DrQv2 [66], the current state-of-the-art policy gradient algorithm on this benchmark, which employs a deterministic actor and hand-tuned exploration.
We also compare with additional baselines that, like SSPG, are based on MaxEnt RL: DrQ [67], CURL [68], and a convolutional version of SAC [39].
Figure 3 (Bottom) shows the performance results after 1.5M environment steps.
DMC tasks yield returns scaled within a set range, [0 , 1000] , which we use for normalization.
Remarkably, also in this domain, SSPG attains state-of-the-art performance with statistically significant improvements over all baselines.
Unlike for the Mujoco tasks, the other considered algorithms based on MaxEnt
RL underperform as compared to the deterministic DrQv2, a result Yarats et al.
[66] attributed to ineffective exploration.
In contrast, SSPG yields performance gains especially on sparser reward tasks where the other baselines struggle (see App. D).
These results validate the scalability of our framework to high-dimensional inputs and its ability to successfully complement MaxEnt RL.
4.2 Properties of serial Markov chain reasoning
We test if theorized benefits of our framework (Sec. 3.5) hold in practical settings with deep networks and stochastic optimization.
We provide further ablation studies and analysis of SSPG in App.
E.
Policy expressiveness.
First, we test the expressiveness of the behavior learned with SSPG using a
Gaussian BT-policy.
We design a series of single-step toy RL problems where the agent needs to position itself on a small 2D environment with a reward function based on unknown goal locations, which we name positional bandits (see App.
C for details).
The objective of these experiments is to isolate how our framework compares with traditional policies for MaxEnt RL to explore the environments and learn to match the true canonical distributions of returns.
As displayed in Fig.
4 A, even in highly multi-modal positional bandits, the SS-policy successfully learns to visit all relevant goals with similar frequencies.
Furthermore, quantizing the state space around the goals reveals that the relative RMC intuitively learns to transition between action-beliefs that visit the different goals as reasoning progresses, with a transition matrix matching a cyclic permutation (App. F).
In comparison, a squashed Gaussian policy expectedly fails to capture the complexity of the canonical distribution, with samples either collapsing to a single mode or covering large suboptimal parts of the action space.
We also show results for a policy based on normalizing flows [69, 70], modeled with a deep expressive network (App. C).
After several attempts, we find these models require orders of magnitude more training iterations and data to learn any behavior that is more complex than a uni-modal distribution.
Yet, even after increasing training by a factor of 1000, we still observe the flow policy distribution collapsing in the more complex positional bandits.
We attribute our findings to training inefficiencies from a lack of proper inductive biases for flow models in the non i.i.d.
RL problem setting [71].
In particular, as flows can assign arbitrarily low probability mass to some regions of the action space, initial local optima can greatly hinder future exploration, exacerbating coverage of the data buffer distribution in a vicious circle.
Policy adaptivity.
Second, we examine the adaptivity of our framework for tackling decisionmaking problems with different complexities.
We compare the average number of reasoning steps ( ¯ N ) performed by SSPG for each task from Sec.
4.1 (Fig.
4 B).
We identify a general correlation between task difficulty and reasoning computation, with complex robotic manipulation and humanoid locomotion problems requiring the most steps.
By concentrating on two representative tasks, we validate the effectiveness of the reasoning process and our adaptive convergence detection strategy with an ablation study where we train SSPG using a fixed number of reasoning steps N fix ∈ { 1 , ⌈ ¯ N ⌉ , ⌈ 3 ¯ N ⌉} .
For the case N fix = 1 , which closely resembles traditional RL, we use double the
UTD ratio to improve performance and offset any training-time gains from multi-step reasoning.
As shown in Fig.
4 C, increasing N fix yields clear performance improvements, validating that agents can greatly benefit from performing longer reasoning processes.
Furthermore, our adaptive SSPG attains the same performance as N fix = ⌈ 3 ¯ N ⌉ and visibly outperforms N fix =
⌈ ¯ N ⌉ .
These results show how different action selection problems require different amounts of reasoning computation and validate the practical effectiveness of our adaptive strategy to detect steady-state convergence.
We obtain analogous findings for additional tasks and values of N fix in App.
F.
Solution reuse.
Last, we examine the effects of the short-term action memory buffer ( ˆ A ) to sample initial action beliefs ( a 0 ) in two tasks.
We evaluate ablating ˆ A , randomly re-initializing a 0 from a uniform distribution.
While there are only minor differences performance-wise between the two approaches (App. F), sampling a 0 from the short-term action memory considerably decreases the number of reasoning steps for convergence (Fig. 4 D).
Moreover, we observe the gap in reasoning efficiency expands throughout training as the agent’s steady-state behavior further improves for the target task.
This result validates that a simple temporal heuristic can provide considerable efficiency benefits, amortizing the additional computational cost of our powerful new framework.
Related work
There have been several prior attempts to extend ubiquitous Gaussian policies [13, 39, 72, 73] with simple normalizing flows [69, 70], both to improve expressiveness [74, 75] and to instantiate behavior hierarchies [76].
Yet, the expressiveness of normalizing flows is coupled with some training challenges [71], which we show can lead to premature convergence to suboptimal solutions in RL (Sec. 4.2).
Other works also considered entirely replacing policy models with gradient-free [4] or gradient-based optimization over the predicted values [77].
Marino et al.
[62] similarly considered learning an optimizer to infer Gaussian behavior [28] with iterative amortization [63].
However, while all these works consider alternative modeling of agent behavior, they are still based on the traditional RL framework of representing decision-making as the output of a fixed process.
Instead, our work entails a conceptually different approach and enables implicit modeling of agent behavior as the result of an adaptive reasoning process, orthogonally providing agents also with additional flexibility to scale computation based on the properties of each individual input state.
Outside RL, there have been efforts to model generation processes with parameterized Markov chains learned to revert fixed noise injection processes acting on data [78–82].
Based on this framework, diffusion models [83–85] recently achieved remarkable results for image generation [85, 86].
While applied to inherently different problem settings, these works share some conceptual resemblances with our framework and highlight the vast scaling potential of implicit modeling.
Conclusion
We introduced serial Markov chain reasoning , a novel framework for modeling agent behavior in
RL with several benefits.
We showed our framework allows an agent to 1) learn arbitrary continuous action distributions, 2) flexibly scale computation based on the complexity of individual actionselection decisions, and 3) re-use prior solutions to accelerate future reasoning.
Hence, we derived
SSPG an off-policy maximum entropy RL algorithm for serial Markov chain reasoning, achieving state-of-the-art performance on two separate continuous control benchmarks.
While for problems with discrete action spaces simple multinomial policy distributions already provide unlimited expressivity, we note that the inherent computational adaptivity of our framework could still yield benefits over traditional fixed policies in these settings.
Furthermore, we believe our motivation and early results provide a strong argument for the future potential of serial Markov chain reasoning, even beyond off-policy RL and simulation tasks.
We provide our implementation for transparency and to facilitate future extensions at sites.google.com/view/serial-mcr/ .
Acknowledgments
We thank Johannes Lutzeyer for providing valuable feedback on an earlier draft of this work.
Edoardo
Cetin would like to acknowledge the support from the Engineering and Physical Sciences Research
Council [EP/R513064/1].
Oya Celiktutan would also like to acknowledge the support from the
LISI Project, funded by the Engineering and Physical Sciences Research Council [EP/V010875/1].
Furthermore, we thank Toyota Motor Europe and Toyota Motor Corporation for providing support towards funding the utilized computational resources.
Appendix A Proofs and extensions A.1 Lemma 3.1, steady-state convergence
For any environment state s , consider a reasoning Markov chain (RMC) defined on a compact action space A with transition probabilities given by π b ( a ′ | a, s ) .
Suppose that inf { π b ( a ′ | a, s ) : a ′ , a ∈
A } > 0 .
Then there exists a steady-state probability distribution function π s ( ·| s ) such that:
[Equation]
Proof.
Let δ = inf { π b ( a ′ | a, s ) : a ′ , a ∈ A } > 0 and let m be the Lebesgue measure .
Since A is compact we have that m
( A ) =
| A | is finite.
Thus, let v be the uniform probability measure on A , we have that:
[Equation]
Hence we have that the resulting RMC is ergodic, irreducible, and aperiodic.
Furthermore, this shows that the entire action space is a small set [87, 88] of the reasoning Markov chain.
This property directly implies the existence of a unique π s and quantifies an exponential bound on the rate of convergence [44, 89]:
[Equation]
where || π bn ( a | a 0 , s ) − π s ( a | s ) || TV represents the total variation distance between the n-step transition probabilities of the RMC and the SS-policy.
We refer to Nummelin [90] for a full formal derivation.
Lemma 3.1 clearly follows from the above results.
A.2 Theorem 3.2, steady-state policy gradient
Let π bθ ( ·| a, s ) be a parameterized belief transition policy which defines a reasoning Markov chain with a stationary distribution given by the steady-state policy π sθ ( ·| s ) .
Let Q s be a real function defined on S × A , with a family of n-step extensions { Q sn } as defined in Eq. 8.
Suppose π b , Q s and their gradient with respect to the parameters θ (denoted ∇ θ ) are continuous and bounded functions.
Then:
[Equation]
Proof.
We start by using the invariance of the RMC’s steady state distribution probabilities, π sθ ( a | s ) , when performing a reasoning step with transition probabilities from the BT-policy π b θ ( a ′ | a, s ) .
Hence, we can decompose ∇ θ E a ∼ π sθ ( ·| s )
[ Q s ( s, a )] in the sum of two distinct terms by applying the product rule:
[Equation]
Leibniz integral rule
[Equation]
From the definition of the n-step extensions being a local approximation of the RMC with no dependence form θ (since ∇ θ Q sn ( s, a ) = 0 ), we can rewrite term (2) as:
[Equation]
Hence, from the relationship between ∇ θ E a ∼ π sθ ( ·| s ) [ Q s 0 ( s, a )] (i
[Equation]
e., ∇ θ E a ∼ π sθ ( ·| s ) [ Q s ( s, a )] ) with (1) and (2) we see that:
[Equation]
Since each n-step extension Q sn has the same exact recursive relationship with the subsequent Q sn + n-step extension, more generally, we have that:
[Equation]
Thus, we can apply Equation 19 recursively to (2) and all resulting ∇ θ E a ∼ π sθ [ Q sn
[Equation]
s
[Equation]
)] terms:
[Equation]
Moreover, from the definitions of the n-step extensions Q sn and the SS-policy π sθ
[Equation]
we see that:
[Equation]
Leibniz integral rule
[Equation]
Leibniz
[Equation]
Therefore, we can use the identity from 21 to simplify Equation 20, and we are left with:
[Equation]
Convergence and finite approximations.
We can use the results derived in Lemma 3.1 (see A.1) to establish some of the properties of the infinite series introduced in Theorem 3.2.
To simplify notation, we will denote each term in the series from Equation 9 with g n ( s, a ) , i.e.:
[Equation]
such that ∇ θ E
a
[Equation]
E
[Equation]
lim
[Equation]
We rewrite the value of each g n ( s, a ) term with the score function , using the identity ∇
[Equation]
log
[Equation]
Leibniz integral rule
[Equation]
Then, we explicitly subtract a baseline, E a ′′ ∼ π sθ [ Q s ( s, a ′′ )] , from the n-step Q-function value Q sn ( s, a ′ ) multiplying the score.
Note, that E a ′′ ∼ π sθ [ Q s ( s, a ′′ )] corresponds to the value function of the SS-policy, which is a baseline independent of a ′ .
Therefore,
E
[Equation]
∇
[Equation]
log
[Equation]
s ))
E
[Equation]
∼
[Equation]
sθ [ Q s
[Equation]
s,
[Equation]
)] (cid:3) = 0 , showing that this subtraction leaves all the gradient terms unchanged [40]:
[Equation]
From our boundedness assumptions
[Equation]
we define existing positive real values Q +
[Equation]
S + such that:
[Equation]
Thus
[Equation]
by noting the relationship
[Equation] [Equation]
we use the bound defined in Equation 18 from Lemma 3.1 to show that:
[Equation]
From Equation 24, it is clear that each term in the series from Theorem 3.2 will converge exponentially fast .
This shows we can estimate E a ∼ π sθ (cid:104)
lim
N →∞ (cid:80)
Nn =
g n ( s, a ) (cid:105) with any arbitrarily small error using a finite N .
Moreover, we can see that each gradient term in Equation 22 is based on the expected deviation of the n-step transition function from the steady-state distribution of the RMC.
This property provides concrete intuition for the proposed adaptive truncation strategy, which considers the expected number of steps before reaching approximate convergence (i.e., π bN ≈ π sθ ).
A similar gradient estimator for the steady-state distribution was also derived with more formal notation in prior works [91].
Heidergott et al.
[92] further extended similar convergence results to broader classes of
Markov chains and even certain unbounded performance functions, providing further motivation for our practical method.
A.3 Policy gradient estimation for arbitrary regularized objectives
We can extend Theorem 3.2 to estimate the policy gradient with respect to a wider class of objectives that involve the expectation over regularized functions :
[Equation]
where the values of Q rθ ( s, a ) might depend on the parameters of the BT-policy, π bθ .
We assume Q rθ and its gradients respect the same regularity assumptions of continuity and boundedness stated in
Theorem 3.2.
The MaxEnt objective from Equation 6 is a particular instance of this setting with the value of Q rθ being dependent on the policy’s entropy, i.e., Q rθ ( s, a )
=
Q s ( s, a ) − α log( π sθ ( a | s )) .
Analogously to the unregularized case, we define a set of n-step extensions for Q rθ for all n = 0 , 1 , 2 , . .
[Equation]
For any state s , we can then apply the product rule to rewrite the gradient of Equation 25 as a sum of two expectations:
[Equation]
Leibniz integral rule
[Equation]
Here, we used the definition of each Q rn ( s, a ) term being a local approximation and having no dependency on the BT-policy to rewrite term (1) as an expectation.
Hence, we can directly apply the unregularized version of Theorem 3.2 to term (1), yielding:
[Equation]
Since E a ∼ π sθ [ ∇ θ Q rθ ( s, a )] =
E a ∼ π sθ (cid:104) E a ′ ∼ π bθ
[Equation]
| a
[Equation]
s )
[ ∇ θ Q rθ
[Equation]
s, a ′ )] (cid:105) , we can merge back the two expectations, obtaining:
[Equation]
The resulting form of the policy gradient generalizes the original Equation 9 from Theorem 3.2, where now the first term in the infinite series explicitly considers the full gradients from the regularized objective.
Hence, applying truncation and the reparameterization trick (as in Equation 13) now yields: ∇ θ
[Equation]
( θ
[Equation]
ϵ
[Equation]
where a ∼ π sθ ( ·| s ) , a 0 = f b ( a, s, ϵ 0 ) , and a i =
f
b ( a i − 1 , s
[Equation]
ϵ i ) for all i =
[Equation] [Equation]
n.
[Equation]
Practical considerations for MaxEnt RL.
The reparameterized extension in Equation 28 is directly analogous to the extension of DDPG policy gradients [12] in SAC [13, 39] when considering the
MaxEnt regularized objective:
[Equation]
To estimate the entropy of the SS-policy, we use the approximation π
[Equation]
a |
[Equation]
n
[Equation]
s ) , which results in a nested Monte-carlo estimator [93] for evaluating this component of the policy gradient with reparameterization.
We find that using N = ⌈ ˆ N ⌉ (Section 3.4) ensures enough samples for practical effectiveness, without necessitating additional debiasing tricks (e.g., [94]).
A.4 Unlimited expressivity
The unlimited representation power of Serial Markov Chain Reasoning comes from the fact that the distribution of agent behavior given by the SS-policy, π sθ , is a mixture model with potentially infinitely many components.
Hence, even a simple Gaussian parameterization of the BT-policy π bθ makes such distribution an arbitrary mixture of Gaussian distributions, which enables agent behavior to approximate any action distribution to arbitrary precision [53].
This is due to Gaussian mixtures being universal approximators of densities [95].
Recent works also show that the same property extends to arbitrary mixture models, with implications for future non-Gaussian extensions of our framework [96].
[Figure]
As described in Section 3.4, we detect convergence in the reasoning process using the multivariate version of the seminal Gelman-Rubin (GR) convergence diagnostic [50, 51].
Given a set of M parallel reasoning chains of length N , a 1: N , this diagnostic computes the pseudo scale reduction factor (PSRF), a score that should expectedly converge to 1 as N increases and the distribution of the different chains approaches the steady-state distribution.
We provide an example in Figure 5 Left where we display how the average PSRF evolves after performing 2-16 reasoning steps after training an agent for the Ant-v2 Mujoco task.
As described in the main text, the PSRF ( R p ) is based on the largest eigenvalue from the matrix product of the sample covariance within ( W ) each of the parallel chains ( (cid:2) a i 1 , . .
. , a iM (cid:3) , for i = 1 , 2 , . .
. , M ) and an unbiased estimate of the true covariance of the steady-state distribution, utilizing the sample covariance between (B) the different chains:
[Equation]
This score can be interpreted as the largest variability in any direction between these two estimates, which can be much higher than the ratio of individual variances computed by the univariate GR scores [52].
Intuitively, as the variance within each chain approaches the variance between independent chains, we can take this as a strong indication that we sufficiently explored the Markov chain’s state-space within each chain.
For our problem settings, we found the PSRF to work well with the default convergence threshold R p < 1 . 1 .
After detecting convergence, we use the number of performed reasoning steps ( N ) to update a running mean ˆ N =
ρ ˆ N + (1 − ρ ) N , as shown in Algorithm 1.
To perform this update, we found a ‘untuned’ choice of ρ = 0 .
99 to also work appropriately.
We observed that for any individual task and training iteration, the majority of the agent’s action selection decisions requiring a number of steps close to ˆ N , as shown by the example in
Figure 5 Right.
Hence, rather than computing the PSRF at every step we found it more efficient to compute the first score only after simulating the RMC for N =
⌊ ˆ N ⌋ reasoning steps.
In particular, starting from a 1 : ⌊ ˆN ⌋ , we consider two cases: If R p ≥ 1 . 1 , we simulate and compute the PSRF for all further steps until convergence.
Otherwise, if R p < 1 . 1 , we backtrack to find the minimum N , such that the subsequence a 1 : N still satisfies our threshold.
With this strategy, we effectively limit the number of evaluations of the PSRF and bolster overall efficiency.
Considerations.
While the assumptions underlying the properties of the GR convergence diagnostic [50, 51] are hard to ensure in practical scenarios, we believe its widespread adoption and empirical effectiveness are strong motivators for its use.
Recent extensions [52] incorporated modern techniques for variance/covariance estimation together with less conservative termination criteria.
Such methods could provide further efficiency gains to our framework, but might also introduce alternative forms of bias while acting and learning.
We leave such exploration and analysis to future work.
C Implementations details
In this Section, we provide descriptions of the SSPG implementations and experimental setups, together with comprehensive hyper-parameters lists.
Since serial Markov chain reasoning is a novel framework, where action-selection is an adaptive iterative process, it adds additional implementation complexity compared to traditional reinforcement learning.
Hence, we share our implementation to ensure reproducibility and facilitate future extensions/comparisons.
By default, the only main requirement of implementing the serial Markov chain reasoning framework is to substitute traditional policies π ( ·| s ) with BT-policies π b ( ·| a, s ) , which take an additional input from the agent’s action space to perform reasoning .
In the considered implementations, we also store a short-term action memory buffer , ˆ A , for acting in the environment, as described in Section 3.3.
Overall, we mostly re-use hyper-parameters and practices from existing policy gradient implementations without major
SSPG-specific tuning.
In total, our algorithm only introduces 4 new main hyper-parameters: 1.
The number of initial action beliefs, M . 2.
The pseudo scale reduction factor (PSRF) threshold to detect approximate convergence.
The size of the short-term memory action buffer, ˆ A . 4.
The coefficient for updating the running mean of performed reasoning steps, ρ (i.e., to update the value of ˆ N for learning).
The values for these new hyper-parameters were mostly chosen based on default values from similar practices and preliminary experimentation.
Overall, SSPG’s performance appeared to be quite robust to a range of reasonable choices and we kept them fixed for all tested domains.
We provide ablation studies and examine performance with different settings in Appendix
E. These results further confirm the claimed robustness.
C.1 OpenAI Gym Mujoco
[Table]
Target critic polyak coefficient 0.
Hidden dimensionality
Fully-connected hidden layers
Nonlinearity
ReLU
Initial entropy coefficient α
Entropy coefficient learning rate 0.
Entropy coefficient β0.
Policy target entropy
Hopper : -1, HalfCheetah : -3,
Walker2d : -3, Ant : -4, Humanoid : -
Critic ensemble size
Critic penalization (uncertainty regularizer [98]) 0.
SSPG-specific hyper-parameters
Number initial action beliefs, M
PSRF threshold 1.
Short-term action memory size
Reasoning steps moving average coefficient, ρ 0.
Our Mujoco SSPG implementation follows common practices employed in the considered baselines and recent literature.
Following REDQ [60] and MBPO [61], we employ a critic ensemble ofmodels and use the suggested task-specific target entropy values for automatic tuning of the MaxEnt coefficient, α [39].
We parameterize each critic with a modern architecture [99] employing 3 fully connected hidden layers as in [98].
Inline with observations for IAPO [62], we find that performing iterative computations with the policy model can easily lead to a value overestimation phenomenon.
Hence, to compute the Q-function’s targets, we use the uncertainty regularizer [98] with a slightlyincreased fixed coefficient of 0.75 (the default value of 0.5 is equivalent to the penalization from clipped double Q-learning [100]).
See Table 1 for a full list of hyper-parameters
[Table]
C.2 DeepMind Control
[Table]
[Table] [Table]
Our DeepMind Control pixel-based SSPG implementation mostly follows the exact hyper-parameters and model specifications from DrQv2 [66], but re-introduces MaxEnt RL in the policy gradient objectives.
The policy target entropy is annealed from 0 .
5 × dim ( A ) to − 1 × dim ( A ) in the first 500K steps, which was done to mimic the exploration standard deviation annealing (from 1 to 0.2 in the first 500K steps) again from the DrQv2 implementation.
Both the other MaxEnt RL parameters and SSPG-specific parameters are kept identical to our experiments on Mujoco tasks.
See Table 2 for a full list of hyper-parameters.
C.3 Positional bandits
To provide visualizations and intuition for the behavior of RL agents adopting the serial Markov chain reasoning framework, we design simple few-dimensional toy tasks , which we call positional bandits .
We used a 1-dimensional positional bandit for the example in Figure 2 and 2-dimensional positional bandits for our expressiveness experiments in Section 4.2.
Positional bandits are defined by an arbitrary list of goal coordinates within a bounded state-space centered around 0.
The action space represents position coordinates, determining which part of the state-space the agent will immediately visit.
The reward of each action is proportional to the distance it brings the agent to the closest goal coordinate.
Thus, in each single-step episode , the agent receives information about its proximity to only one of the goals.
In these environments, we train all agents with the MaxEnt objective, making optimal behavior correspond to visiting all-goals with similar frequencies.
These minimal problems aim to isolate and assess the ability of different agents of modeling arbitrary multi-modal distributions of behavior in the non i.i.d.
RL problem setting which also involves a non-trivial exploration challenge.
We evaluate light-weight versions of SSPG, SAC, and a flow-based version of SAC.
We largely follow the specifications from the Mujoco experiments, with a shared UTD of 1, a single critic, 50 random exploration steps, a fixed α , and with all models (except the flow-based policy) using 2-layer fully-connected architectures with 32 hidden dimensions.
We implemented the flow-based version of SAC, based on inverse autoregressive flows [101], as also considered in related RL works [62, 75].
In particular, we apply to the policy model’s outputs two additional flow transformations, each parameterized by a 2-layer fully connected network.
We found that this class of models benefits from wider policy networks, thus, we increased back the hidden dimensionality to 256 (as in our main Mujoco experiments).
We found 1000 steps of experience collection/network updates was sufficient for behavior to convergence in SSPG and standard SAC.
Instead, the flow-based version of
SAC required much additional experience/training before attaining behavior resembling any actual multi-modality, hence, we considered its results after 1M steps.
C.4 Other considerations
In this work, we implemented SSPG with minimal extensions to existing policy gradient models and parameters.
The purpose of this choice was to remove confounding factors related to hyper-parameter and model tuning, and evaluate our framework mainly for its conceptual properties over traditional
RL agents.
For the same reason, we concentrated our empirical analysis on established evaluation benchmarks and toy problems, rather than exploring new challenging applications.
This leaves much potential for future work to investigate new model architectures, training practices, and push the limit of serial Markov chain reasoning .
Early versions of our framework considered using some recurrent memory components as input to the
BT-policy.
However, to preserve the Markov property, this modification would require considering an extended state space for the RMC, consisting of both action-beliefs and memory hidden states.
Preliminary results showed that such extended state space makes convergence require an increased amount of reasoning steps and optimization of the BT-policy more unstable.
D Extended evaluation results
In this Section
[Table]
we provide granular per-task results and further aggregate evaluations with the Rliable evaluation protocol [55] after collecting different numbers of steps.
We also compare the performance of SSPG with a new baseline integrating prior state-of-the-art algorithms with the normalizing flow model introduced in the first part of Section 4.2.
Additionally, we report the mean running times for the considered algorithms and the main baselines we run to obtain performance results.
Together with details about our hardware setup, these should give a solid intuition of the relative computational training cost of each algorithm.
Furthermore, we also provide results for the performance of SSPG when ‘clipping’ the maximum number of reasoning steps allowed for each action-selection
[Table]
and the relative evaluation times to analyze potential deployment-time efficiency and trade-offs.
D.1 OpenAI Gym Mujoco
[Table]
Per-task results.
In Table 3 we provide the per-task results collected at 100K and 200K steps for the considered OpenAI Gym Mujoco tasks.
For both experience thresholds, SSPG obtains the best average performance in 5/6 tasks, and still lags very close the model-based MBPO [61] algorithm for the remaining task (HalfCheetah-v2).
SSPG converges much earlier than other algorithms, even while performing many less optimization steps (REDQ, REDQ-FLOW, MBPO, and SAC-20 all employ a
UTD of 20, while we train SSPG with a UTD of 10, see Section 4).
Extended aggregate results.
In Figure 6, we provide additional aggregate metrics collected at 50K and 200K steps, using the same statistical tools described in Section 4.
We see most of the considerations for the evaluation at 100K steps, equally hold at these different experience regimes.
Most notably, SSPG aggregately outperforms all prior algorithms with statistically meaningful gains, for the Neyman-Pearson statistical testing criterion [64] and also exhibits stochastic dominance [65] after 200K training steps.
D.2 DeepMind Control
[Table]
Per-task results.
In Table 4 we provide the per-task results collected at 1.5M and 3M steps for the considered DeepMind Control tasks from pixel observations from [66].
SSPG obtains the best average performance in 10/12 and 9/12 tasks, respectively.
We can observe most notable gains in some of the tasks that pose harder exploration challenges (e.g., cartpole_swingup_sparse) and especially in lower-data regimes (e.g., finger_turn_hard).
We believe this could be an indication that serial Markov chain reasoning complements particularly well the MaxEnt reinforcement learning framework and is able to overcome some of its observed limitations for tackling sparse reward, pixel-based tasks [66].
Extended aggregate results.
In Figure 7, we provide additional aggregate metrics collected at 500K and 3M steps, using the same statistical tools described in Section 4.
Again, these results are consistent with the evaluation at 1.5M steps.
In particular, we highlight ubiquitous statistical significance and stochastic dominance [65] after training from 500K environment steps.
D.3 Comparison with normalizing flows
We provide results comparing SSPG with additional new baselines, extending state-of-the-art algorithms based on the traditional RL framework with the normalizing flow policy model considered in
Section 4.2.
To counteract the exploration collapse phenomenon observed in the positional bandits experiments, we found it beneficial to add a small amount of fixed Gaussian noise on top of the already-stochastic policy when collecting training data ( σ = 0 . 05 ).
While powerful flow models can also provide unlimited expressiveness, they are still based on the traditional RL framework and represent decision-making as the output of a fixed process.
Thus, these experiments can provide insights regarding the contribution of adaptivity to our SSPG framework.
We compare the relative improvements from SSPG and flows with respect to the original performance of the considered state-of-the-art baselines in either the OpenAI Gym Mujoco and DeepMind Control benchmarks
[Table]
described in
[Table]
In Table 5 we provide the per-task comparisons.
incorporating flows with REDQ appears to provide limited and inconsistent performance benefits as compared to incorporating our serial Markov chain reasoning framework.
This is in line with our intuition justifying SSPG’s superiority and with results reported in prior off-policy reinforcement learning work which show flows provide marginal gains over standard Gaussian or deterministic policies (e
[Table]
g., [62, 75])
[Table] [Table]
In Table 5 we provide the per-task comparisons.
Similarly to our results in the OpenAI Gym Mujoco tasks
[Table]
incorporating flows with DrQv2 also appears to provide limited and inconsistent benefits as compared to the benefits of SSPG.
We note that in both benchmarks
[Table]
flows appear to particularly help earlier in training (100K steps and 1.5M steps, respectively), likely due to some initial exploration benefits.
However, its
In Figure
[Table]
we provide aggregate metrics for the normalizing flows agents, using the same statistical tools described in Section 4.
We considered performance after 100K and 1.5M steps for the OpenAI
Gym Mujoco and DeepMind Control benchmarks, since its where normalizing flows appear to yield the most benefits.
However, even in this experience regime, incorporating flows does not provide statistically significant gains and does not exhibit stochastic dominance [65] as SSPG.
D.4 Computational requirements
[Table] [Table]
We record the average training time from executing the different considered algorithms for each training epoch , consisting of 1000 environment steps for OpenAI Gym and 10000 environment steps for DeepMind Control.
The most relevant specifications of our hardware setup are an NVIDIA RTX 3090 GPU and an AMD Ryzen Threadripper 3970x CPU.
We based our implementations on the shared code from [98] and [66] for OpenAI Gym and DeepMind Control tasks, respectively.
Training times.
We show the average training times in Table 7.
For the OpenAI Gym Mujoco tasks, the increased UTD ratio severely affects computation requirements.
Thus, by halving the UTD ratio of SSPG, we are able to bring its epoch training time lower than SAC-20 and all other considered baselines.
However, we did not explore lowering the number of updates for our experiments in
DeepMind Control, as DrQv2 already employs a low UTD ratio of 0.5.
Thus, DrQv2 is about 30% computationally faster than SSPG.
We believe the extra computation for carrying out the reasoning process is one of the main limitations of the serial Markov chain reasoning framework.
However, for real-world problems, physical constraints related to environment feedback are often the main bottlenecks for the overall time efficiency.
In such regard, we argue that the considerable sampleefficiency improvements from our novel framework should be much more relevant than increased computational cost.
D.5 Analysis of the cost of acting during deployment
[Table]
We report the performance of SSPG when ‘clipping’ the maximum number of reasoning steps allowed for each action-selection to fixed values, during evaluation only.
We evaluate final checkpoints of agents learned without any clipping by our unmodified SSPG.
We consider a subset of four total environments from both OpenAI Gym Mujoco (i
[Table]
e., Humanoid-v2, and Ant-v2) and DeepMind
Control (i
[Table]
e., cheetah_run, and quadruped_run) for which SSPG displays different average reasoning requirements.
As shown in Table 8, clipping to as low as four reasoning steps marginally affects the performance of SSPG, which still always surpasses the scores achieved by standard reinforcement learning baselines.
SSPG is less affected by this form of deployment-only clipping than by fixing the number of reasoning steps for both training and evaluation phases.
We motivate this finding by noting that better capturing the canonical distribution of returns from the critic by performing additional reasoning steps has also significant benefits during the training phase of off-policy algorithms.
In particular, this allows the agent to achieve better exploration and more easily correct the critic in the areas of the action space where its predictions are erroneously optimistic.
When training with the unmodified SSPG algorithm, this benefit is fully retained, justifying the performance superiority than our previous fixed-steps ablations.
[Table]
We report the average rollout time during deployment of each of our implementations.
As shown in Table 9, using SSPG does increase the average rollout time over standard reinforcement learning baselines.
However, the additional time required for action-selection scales sub-linearly with the number of reasoning steps, and appears to be dominated by other fixed costs, such as simulating the environment and converting observations to tensor objects.
This is in contrast with the other more expensive iterative baseline, IAPO [62], which performs gradient-based optimization at each acting step.
Differences with standard reinforcement learning are even more marginal in the visual
DeepMind Control environments, where the most expensive part of the computation is from encoding the RGB input observation with a convolutional encoder (which needs to occur only once before performing the reasoning steps with the policy ‘head’).
Moreover, we note that in many real-world applications, the additional acting cost would be greatly inferior to the actuation time costs when using distributed hardware.
However, clipping the number of reasoning steps still remains a valuable option, as examined in the experiment above.
E Ablations and parameter studies
We perform additional experiments to evaluate SSPG’s design choices and test alternative optimization setups for serial Markov chain reasoning .
We focus on the Humanoid-v2 and quadruped_run tasks, which we found to be generally representative of overall agent behavior in the OpenAI Gym Mujoco and DeepMind Control domains, respectively.
For each experiment, we report the performance curves for average returns (higher is better) and number of reasoning steps (lower is better).
E.1 Loss backpropagation
Following the policy gradient estimation derived in Section 3.2, we optimize the BT-policy by recording its gradients with respect to the following operations: 1) We perform one (reparameterized) reasoning step to output an action-belief, a 0 ∼ π bθ ( s, a ) .
We simulate a local approximation of the (reparameterized) RMC for N steps, yielding a 1: N . 3)
We compute and backpropagate the outputs of all Q sϕ ( s, a n ) , for n = 0 , ..., N .
A practical alternative would be to carry out step (2) by performing N reasoning steps with the (reparameterized) BT-policy itself, rather than with a local approximation of the RMC.
While this change would not affect the objective value from each
Q sϕ ( s, a n ) , it would make the computation record additional gradient dependencies to the BT-policy’s parameters from each reasoning step , rather than just the first.
As we already have to perform forward and backward passes for the N actions computed in step (2), this extension would not come at a significant computational overhead and would be reminiscent of differentiation through a traditional recurrent model.
In practice, we expect this change would enable to reduce optimization variance, as we are optimizing the direct influence that the BT-policy has on the RMC starting from N different initial action-beliefs for each sample.
Yet, this alternative optimization procedure is not supported by our theoretical intuition and would likely introduce further bias into the policy gradient estimation.
We show experimental results in Figure 9.
We evaluate both a truncated version (truncated full backprop.) and an un-truncated version (full backprop.)
that perform full backpropagation through the RMC when optimizing SSPG.
The truncated version still optimizes the BT-policy with respect to the full N reasoning steps, but prevents the gradients of future steps backpropagating to earlier action-beliefs (somewhat reminiscent of truncation in RNN optimization).
Surprisingly, we observe no major differences performance-wise.
However, we do observe somewhat higher variance between different experiments when performing full backpropagation and slightly lower final returns, likely symptomatic of the additional bias.
On the other hand, the un-truncated version does experience marginally faster initial learning in quadruped_run, indicating there could be some practical benefits from backpropagating multiple reasoning steps to the BT-policy’s parameters.
We leave exploration of such potential extensions for future work.
As described in Section 3.3, we perform the reasoning process from a set of initial action-beliefs a 0 .
The size of a 0 (denoted | a 0 | ) influences the number of initial starting modes and the accuracy of the empirical PSRF computation for convergence detection.
By default SSPG employs | a 0 | = 64 .
As shown in Figure 10, doubling or halving | a 0 | affects convergence speed, with a greater number of initial action-beliefs leading to a slightly faster reasoning process.
We attribute this phenomenon to the conservativeness of the proposed PSRF convergence rule [52], which becomes harder to satisfy when evaluated from fewer parallel chains.
However, these results do not take into account the cost of performing each reasoning step which is itself affected by | a 0 | .
SSPG’s overall efficiency should be quite robust to most choices for this hyper-parameter with hardware optimized for parallel computation.
E.3 Short term action memory buffer size
As described in Section 3.3, we sample initial action-beliefs for reasoning from the short-term action memory buffer, ˆ A .
The size of ˆ A (denoted | ˆ A | ) determines how far into the past to store action-beliefs that can facilitate future reasoning.
In practice, we found that the simple heuristic of setting | ˆ A | = | a 0 | works well.
In Figure 11, we examine the effects of doubling the buffer size from 64 to 128.
Again, we observe this change leads to an increase in the expected reasoning steps before reaching convergence.
This result indicates that the most recent past actions are generally most relevant, which we believe is a consequence of the fine temporal discretization of the considered environments.
We compute the quantized transition probabilities of the RMCs in the considered positional bandits from our Policy expressiveness analysis (Section 4.2).
For each positional bandit we randomly sample a set of 1000 action-beliefs within a radius around each goal and compute the likelihood of transitioning between any two such action-beliefs.
We sum and normalize these probabilities both with respect to the input and output action beliefs within each goal.
Hence, we obtain a transition matrix for the discretized RMC, showing the probability of updating the current action-belief based on the resulting closest goal.
We find that the BT-policy intuitively learns to consider action-beliefs to reach all different goals in turn.
For instance, in the positional bandit with three goals, performing a reasoning step with the learned BT-policy from an action-belief that would land nearby goalwould almost certainly lead to an action-belief that would land near goal 3.
Similarly, an action-belief landing near goal 2 would follow from an action-belief that would land near goal 3, and would almost certainly lead to an action-belief that would land near goal 1 (with less than 10 − 6 probability of deviating).
This behavior is summarized from visualizing the relative transition matrices following the described quantization in Figure 12.
F.2 Reasoning with fixed action beliefs
We provide further results in support of our Policy adaptivity analysis by comparing the performance of our adaptive SSPG with performing alternative fixed numbers of reasoning steps in 4 different tasks with diverse average reasoning costs (see Figure 4 B).
As shown in Figure 13 and also Figure
C, increasing the number of reasoning steps monotonically improves performance.
Yet, we did not observe any case where our adaptive strategy underperforms as compared to any fixed number of steps.
We believe this is further empirical evidence of how a framework that can adaptively dedicates different amounts of computation to different action-selection problem can provide great efficiency and scaling benefits.
F.3 Effects of short-term action memory buffer
As stated in the Solution reuse analysis, and shown in Figure 14, re-initializing randomly a 0 hurts efficiency but not performance.
We believe this is an indication that the BT-policy effectively learns to bootstrap information in each previous action-belief while reasoning.
Intuitively, starting from an initial action-belief that contains no information about optimal behavior makes the relative actionselection problem harder.
Yet, SSPG still successfully enables to recover very similar performance via adaptively performing additional reasoning steps to compensate.
G Societal impact
We proposed a new framework for modeling autonomous decision-making in reinforcement learning.
Thus, the societal impact of our work is tightly related to the impact of the reinforcement learning field.
While autonomous systems can offer many benefits, as they become more commonplace, they might introduce ethical issues such as privacy, surveillance, and bias.
If unregulated, automation might also accentuate economic inequalities and have non-trivial environment impact.
Policy making strategies and regulations are increasingly needed to mitigate these risks, which we believe should not compromise the field’s advancements given its countless potential positive implications.