SPP-RL: State Planning Policy Reinforcement Learning

We introduce an algorithm for reinforcement learning , in which the actor plans for the next state provided the current state . To communicate the actor output to the environment we incorporate an inverse dynamics control model and train it using supervised learning . We train the RL agent using off-policy state-of-the-art reinforcement learning algorithms : DDPG , TD3 , and SAC . To guarantee that the target states are physically relevant , the overall learning procedure is formulated as a constrained optimization problem , solved via the classical Lagrangian optimization method . We benchmark the state planning RL approach using a varied set of continuous environments , including standard MuJoCo tasks , safety-gym level 0 environments , and AntPush . In SPP approach the optimal policy is being searched for in the space of state-state mappings , a considerably larger space than the traditional space of state-action mappings . We report that quite surprisingly SPP implementations attain superior performance to vanilla state-of-the-art off-policy RL algorithms in a wide class of robotic locomotion environments . 1 INTRODUCTION . Research on reinforcement learning ( RL ) has brought a tremendous number of successful applications in diverse fields of science and technology . Application areas of RL can be split into two classes : discrete and continuous . Here , we are interested in continuous simulation environments , mostly in the robotics domain . Despite the magnitude of successful applications of RL in this domain , two of the main issues : sample efficiency and interpretability of RL trained agents , persist . These problems have utmost practical importance , especially for mission-critical applications . The current methods often require a vast amount of experience for training . The decision-making process of trained agents is not interpretable , sometimes resulting in finding proxy solutions . Thus , it is vital to research new RL algorithms that may partially solve the mentioned problems . Traditionally , RL is based on the principle of searching for the optimal policy within the space of state-action mappings . We propose a new algorithm based on the principle of training an actor ( a policy ) operating entirely in the state space ( state-state mappings ) . We call such policies the state planning policies ( SPP ) , whose actions determine desired trajectories in the state-space . The task of training SPP may initially seem infeasible due to a significantly larger dimension of states than of actions . Nonetheless , quite surprisingly , we show that the approach is feasible and often leads to significant improvements in average performance for a class of robotic locomotion tasks . We call our approach State Planning Policy Reinforcement Learning ( SPP-RL ) . It is a generic approach for problems specified using continuous environments . The main building block of SPP-RL – the RL agent can be implemented using virtually any RL algorithm . We chose to develop our approach using the state-of-the-art off-policy DDPG ( Lillicrap et al. , 2016 ) , TD3 ( Fujimoto et al. , 2018 ) , and SAC ( Haarnoja et al. , 2018a ) algorithms . Note that , in SPP-RL we need another trainable model to communicate the policy output to the environment ; as such , we incorporate a learnable inverse dynamics control model ( IDM ) , see Fig . 1 . The overall algorithm optimizes the policy simultaneously with CM . To ensure that the policy target states satisfy physical and under-actuation constraints , we formulate a constrained optimization objective for policy training . Our work lies within the category of RL methods that already implemented state-state policies , including work on hierarchical RL like Nachum et al . ( 2018 ) , and the D3G algorithm from Edwards et al . ( 2020 ) . Summary of results . Although the SPP-RL algorithm searches for the optimal policy within a much larger space , our performance benchmarks revealed that SPP-RL implementations often outperform their vanilla RL counterparts . Experiments performed in Ant & Humanoid tasks from MuJoCo suite ( Todorov et al. , 2012 ; Brockman et al. , 2016 ) reveal that SPP-DDPG outperforms DDPG . Experiments in Safety-Gym Level 0 environments ( Ray et al. , 2019 ) demonstrate that SPP-TD3 and SPP-SAC outperform by a great margin TD3 and SAC , respectively . Experiments in AntPush task ( Nachum et al. , 2018 ) show that SPP-TD3 outperforms hierarchical RL method HIRO ( Nachum et al. , 2018 ) and provides some interpretability of the agent behavior . We hypothesize that the superior performance of SPP-RL in the tested continuous environments originates in more efficient state-space exploration by state-state policies than traditional state-action policies ; here noise is being added to target states rather than actions . To argue this , we performed series of experiments , including evaluation of a shadow agent utilizing experience from SPP and vanilla replay buffers ( Sec . 5.4 ) and a study of the distributions of states gathered in different replay buffers ( App . E.2 ) . To analyze which features of the algorithm are crucial for superior performance , we report on a thorough ablation study of SPP-TD3 ( App . E.1 ) . We implemented SPP-RL methods as a modular PyTorch library shared as open-source . SPP-RL algorithms are derived from their vanilla RL counterparts , making extending the library with new RL algorithms straightforward . We also share videos with test episodes of the trained agents to accompany benchmark plots ( web , 2021 ) . Last but not least , we demonstrate theoretical convergence of clipped Double Q-learning implemented in SPP-TD3 algorithm within the classical stochastic processes , finite-state/action spaces setting and in environments based on rigid body dynamics model , refer to App . B.3 . 1.1 RELATED WORK . We present a ( non-exhaustive ) list of related works ; refer to Tab . 1 for a perspective on related work . The closest approach to ours is the D3G algorithm introduced by Edwards et al . ( 2020 ) , which includes state planning policies , and introduces a novel form of the value function defined on statenext state pairs . There are two main ways our method is distinct . First , SPP employs the classical formulation of the value function . Also , we do not include a forward dynamics model nor the cycle loss . Instead , to guarantee consistency of the policy target-states in SPP , we formulate a constrained optimization problem ( compare Fig . 2 ) solved via Lagrangian optimization . Our work builds on the classical RL algorithms going back to REINFORCE ( Williams , 1992 ) , Asynchronous Actor-Critic ( Mnih et al. , 2016 ) , and especially the off-policy actor-critic algorithms including Q-Prop ( Gu et al. , 2017 ) , DDPG ( Lillicrap et al. , 2016 ) , SAC ( Haarnoja et al. , 2018a ; b ) , and TD3 ( Fujimoto et al. , 2018 ) . State planning policies have been used in hierarchical RL ( HRL ) methods like HIRO ( Nachum et al. , 2018 ) and FuN ( Vezhnevets et al. , 2017 ) . Contrary to HRL SPP-RL approach does not employ a hierarchy of multiple policies nor state conditioned value functions . Training predictive models ( like IDMs ) is fundamental for the model-based RL approach including algorithms : a locally linear latent dynamics model ( Watter et al. , 2015 ) , model-based planning for discrete and continuous actions ( Henaff et al. , 2017 ) , model-predictive control ( Chua et al. , 2018 ) , and model based policy optimization ( Janner et al. , 2019 ) . We deployed IDMs for mapping currenttarget states to actions ; other applications of IDMs in RL include the context of planning : search on replay buffer ( Eysenbach et al. , 2019 ) , episodic memory graph ( Yang et al. , 2020 ) , and topological memory for navigation ( Savinov et al. , 2018 ) . Existing many other applications of IDM in context of RL including : policy adaptation during deployment ( Hansen et al. , 2021 ) , sim to real transfer ( Christiano et al. , 2016 ) , adversarial exploration ( Hong et al. , 2020 ) , and curiosity-driven exploration ( Pathak et al. , 2017 ) . 1.2 BACKGROUND . Following the standard setting used in RL literature , we work with infinite horizon Markov decision process ( MDP ) formalism ( S , A , P , r , ρ0 , γ ) , where S is a state space , A is a action space , P : S × A×S → [ 0 , 1 ] is a transition probability distribution , r : S ×A → R is a reward function , ρ0 is an initial state distribution , and γ ∈ ( 0 , 1 ) is a discount factor . From now on we assume that the MDP is fixed . In RL the agent interacts with E in discrete steps by selecting an action at for the state st at time t , causing the state transition st+1 = E ( st , at ) , as a result the agent collects a scalar reward rt+1 ( st , at ) , the return is defined as the sum of discounted future rewardRt = ∑T i=t γ ( i−t ) r ( si , ai ) . The goal in RL is to learn a policy that maximizes the expected return from the start distribution . 2 STATE PLANNING POLICY REINFORCEMENT LEARNING APPROACH . Our SPP approach is rooted in state-state reinforcement learning , by which we mean setting in which RL agent is trained to plan goals in the state-space , the approach already employed e.g. , in HRL , planning , D3G RL algorithms ( see Tab . 1 ) . In SPP a state planning policy π given the current state st outputs zt – the desired target state to be reached by the environment in the next step . Forcing the environment to reach the desired state requires translating the target state to a suitable action at . Hence , we employ an additional model capable of mapping the current state-target state pair ( st , zt ) to the action at – a ( trainable ) IDM model . Ideally , we like to have consistency zt ( st ) ≈ st+1 . The consistency can not be guaranteed a-priori , is rather achieved in SPP setting by employing a constrained optimization approach . A diagram illustrating SPP approach is presented in Fig . 1 . We have freedom of choice of the particular RL algorithm ( RL agent ) and IDMs implementations . Currently , we use feed-forward neural networks , and RL Agent using implementations of the stateof-the-art off-policy RL algorithms : DDPG ( Lillicrap et al. , 2016 ) , TD3 ( Fujimoto et al. , 2018 ) and SAC ( Haarnoja et al. , 2018a ) . We present details of SPP-RL implementation in Sec . 4 using as the example SPP-DDPG . The encountered experiences during the execution of an off-policy RL algorithm are stored in replay buffer D. The main building block of SPP-RL are the state planning policies , intuitively a state planning policy selects a desired trajectory in the state-space of the environment . Definition 1 . We call a state planning policy a map πθ : S → P ( S ) parametrized using a vector of parameters θ ∈ Rn ( θ ) , and we denote πθ ( z|s ) a probability of the desired target state z ∈ S for the given current state s ∈ S. We call a deterministic state planning policy a parametrized map πθ : S → S , and we denote πθ ( s ) = z . We assume that π has continuous and bounded derivatives with respect to θ . We will call state planning policy whenever it is clear from the context deterministic/stochastic and omit the parameter subscript π = πθ . Besides the state planning policy ( Def . 1 ) the second main building block of the overall SPP agent is a model for mapping the current state-target pair ( st , zt ) to suitable action at . Following the existing literature , we call such model the inverse dynamics control model ( IDM ) , or simply the control model . Definition 2 . For a given MDP ( S , A , P , r , ρ0 , γ ) . Let s , z ∈ S. We define the control model : CM : S × S → A , CM ( s , z ) = a , i.e . for the given two states CM computes the action a . We call s , z the initial state and the target state respectively , where a informally satisfies argmaxb∈S P ( s , a , b ) ∼ z for stochastic E , or z ≈ E ( s , a ) for deterministic E. Obviously , in order to work , SPP requires consistency of the target states generated by the policy with the actual next-states of the environment . We call this property the state consistency property ( or simply consistency ) of π , refer to Fig . 2 . As it may be intractable to verify SPP for all possible interactions in continuous environments , we are interested in guaranteeing the state consistency for the experiences stored in the replay buffer . Property 1 . Let D be a replay buffer , CM be an IDM and π be a ( SPP ) policy . We say that π has the state consistency property with threshold d > 0 if it holds that E ( st , st+1 ) ∈D zt∼π ( st ) [ ∥st+1 − zt∥22 ] ≤ d , for deterministic π we have zt = π ( st ) . Given ( zt , st+1 ) , we call distance ∥zt − st+1∥22 the stateconsistency distance ( refer Fig . 2 ) . We will often assume that d is known from context and omit ’ with threshold d > 0 ’ .

Summary: The paper describes an approach where a state-state mapping is learned (called state-planning policy or SPP) coupled with an off-policy RL system (the authors included experiments with DDPG, TD3,  and SAC).  Learning a state-state mapping, Q(s,s'), needs also to learn an inverse dynamic model to estimate which action (a) will take the current state (s) to the next state (s').  The authors framed the problem as a constrained optimization approach, where the objective is to maximize the sum of discounted rewards such that the differences between the target states (produced by the policy) and the actual states (produced by the environment with the action from the control model) are within a certain threshold distance. In particular, the authors solved the optimization problem via the Lagrange multiplier method.  Experiments were performed in different domains and with different state-of-the-art DRL systems.