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 difﬁculty 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-speciﬁc 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 efﬁciency 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 ﬁxed 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 ﬂexible 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 deﬁnes 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 difﬁculty 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 efﬁciency 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], deﬁned 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 2 S , based on which it selects an action from its own action space, a t 2 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 =
-]
Q Tt =0 ⇡ ✓ ( 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 ⇡✓ ( ⌧ ) h P
Tt =
  t r ( s t , a t ) i .
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 efﬁciently 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 0 | 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 deﬁned as:
[Equation]
n
[Equation]
Given a compact action space and a BT-policy with a non-zero inﬁmum 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) deﬁned on a compact action space A with transition probabilities given by ⇡ b ( a 0 | a, s ) .
Suppose that inf { ⇡ b ( a 0 | a, s ) : a 0 , a 2 A } > 0 .
Then there exists a steady-state probability distribution function ⇡ s ( ·| s ) such that:
[Equation]
[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  0 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 deﬁned, 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 7!
R for n = 0 , 1 , 2 , . . . , deﬁned as:
[Equation]
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 ﬁxed and immutable copy of the current ⇡ b ✓ .
We use this deﬁnition 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 deﬁnes a reasoning Markov chain with a stationary distribution given by the steady-state policy ⇡ s ✓ ( ·| s ) .
Let Q s be a real function deﬁned on S ⇥ A , with a family of n-step extensions { Q sn } as deﬁned in Eq. 8.
Suppose ⇡ b , Q s and their gradient with respect to ✓ (denoted r ✓ ) 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 ﬁnite 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 0 , s ) ⇡ ⇡ s ✓ ( a | s ) and, thus, Q sn ( s, a 0 ) ⇡
[Equation]
Therefore, the value of each Q s n ( s, a 0 ) will be independent of the BT-policy’s action a 0 , such that r ✓ E a 0 ⇠ ⇡ b ✓ ( ·| a,s )
[ Q sn ( s, a 0 )] ⇡ 0 .
In other words, each subsequent step in the RMC introduces additional randomness that is independent of a 0 , 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 r 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 deﬁned in Eq. 8).
By recursively applying this relationship (Eq. 11) to r a 1 Q sn  1 ( s, a 1 ) and all subsequent gradient terms we obtain:
[Equation]
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 sn ( s, a 0 ) 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 + 1 while R p > 1 .
1 do a N + 1 ⇠ b ·| a N ) N N
Update R p with : N .
N ⇢ ˆ N + (1  ⇢ ) N .
A ˆ A [ a 1 : N output: a ⇠ a 1 : N input: D , data buffer ( s, a, s 0 , r ) ⇠ D a 0 ⇠ ⇡ b ✓ ( ·| a, s 0 ) for n 0 , d
ˆ N e do Q sn Q s  ( s 0 , a n ) .
Eq. ✏ n +1 ⇠ N (0 , 1) , a n +1 = f b ( a n , s, ✏ n +1 ) r ✓ Q s  r ✓ ( P d N e n =0 Q sn ) .
Thm.
3.2 arg min ✓ J ( ✓ ) . Eq.
6 a⇠ a 1: d ˆ N e 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 ﬁxed sized buffer, called the short-term action memory , ˆ A , and use them to construct a 0 .
We ﬁnd 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 sufﬁcient 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 ﬁxed 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 d ˆ N e as the number of reasoning steps to simulate the RMC with ⇡ b ✓ when computing gradients from Eqs.
d ˆ N e 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 simpliﬁed summary of our adaptive reasoning process for acting and learning in Algs.
3.5 Advantages of serial Markov chain reasoning
Based on the above speciﬁcation, 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 inﬁnitely 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 ﬂexibly spend different amounts of computation time based on the complexity of each action-selection problem, with potential gains in both precision and efﬁciency.
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 speciﬁed in Sec.
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 proﬁles [57], interquantile mean (IQM), and probability of improvements over baselines with the Mann-Whitney U statistic [58].
The reported ranges/shaded regions represent 95% stratiﬁed bootstrap conﬁdence intervals (CIs)
In App.
D, we provide per-task results and further statistical analysis.
For each experiment, we collect the returns of SSPG over ﬁve 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-efﬁciency 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 beneﬁts 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-efﬁciency, only reducing the amortization gap beyond direct estimation might not provide signiﬁcant beneﬁts.
Instead, serial Markov chain reasoning’s improved expressivity and ﬂexibility 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 speciﬁcations 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 signiﬁcant 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 beneﬁts 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 ﬂows [69, 70], modeled with a deep expressive network (App. C).
After several attempts, we ﬁnd 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 ﬂow policy distribution collapsing in the more complex positional bandits.
We attribute our ﬁndings to training inefﬁciencies from a lack of proper inductive biases for ﬂow models in the non i.i.d.
RL problem setting [71].
In particular, as ﬂows 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 difﬁculty 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 ﬁxed number of reasoning steps N fix{ 1 , d ¯ N e , d 3 ¯ N e } .
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 beneﬁt from performing longer reasoning processes.
Furthermore, our adaptive SSPG attains the same performance as N fix = d 3 ¯ N e and visibly outperforms N fix =
d ¯ N e .
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 ﬁndings 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 efﬁciency 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 efﬁciency beneﬁts, 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 ﬂows [69, 70], both to improve expressiveness [74, 75] and to instantiate behavior hierarchies [76].
Yet, the expressiveness of normalizing ﬂows 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 ﬁxed 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 ﬂexibility 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 ﬁxed 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 beneﬁts.
We showed our framework allows an agent to 1) learn arbitrary continuous action distributions, 2) ﬂexibly 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 beneﬁts over traditional ﬁxed 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.
[-Nn =
-]