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: This paper proposes a new RL algorithm whereby the policy selects a new state rather than action, with constraints to ensure the next state selected is a valid one. The contribution is the new algorithm "SPP-RL" which can be applied to off policy algorithms such as TD3/SAC. There is experimental evidence this may be an effective approach in some settings.