Learning powerful policies and better dynamics models by encouraging consistency

1 INTRODUCTION . Reinforcement Learning consists of two fundamental problems : learning and planning . Learning comprises of improving the agent ’ s current policy by interacting with the environment while planning involves improving the policy without interacting with the environment . These problems evolve into the dichotomy of model-free methods ( which primarily rely on learning ) and modelbased methods ( which primarily rely on planning ) . Recently , model-free methods have shown many successes , such as learning to play Atari games with pixel observations ( Mnih et al. , 2015b ; Mnih et al. , 2016 ) and learning complex motion skills from high dimensional inputs ( Schulman et al. , 2015a ; b ) . But their high sample complexity is a still a major criticism of the model-free approaches . In contrast , model-based reinforcement learning methods have been introduced in the literature where the goal is to improve the sample efficiency by learning a dynamics model of the environment . But model-based RL has several caveats . If the policy takes the learner to an unexplored state in the environment , the learner ’ s model could make errors in estimating the environment dynamics , leading to sub-optimal behaviour . This problem is referred to as the model-bias problem ( Deisenroth & Rasmussen , 2011 ) . In order to make prediction about the future , dynamics models are unrolled step by step which leads to the issue of “ compounding errors ” ( Talvitie , 2014 ; Bengio et al. , 2015 ; Lamb et al. , 2016 ) : an error in modeling the environment at time t affects the predicted observations at all subsequent time-steps . This problem is much more challenging for the environments where the agent observes high-dimensional image inputs and not compact state representations . On the other hand , modelfree algorithms are not limited by the accuracy of the model , and therefore can achieve better final performance by trial and error , though at the expense of much higher sample complexity . In the model-based approaches , the dynamics model is usually trained with supervised learning techniques mostly just by observing the data . On the other hand , there ’ s enough evidence that humans learn the environment dynamics not just by observing the environment but also by interacting with the environment ( Cook et al. , 2011 ; Daniels & Nemenman , 2015 ) . This leads to an interesting possibility . The agent could consider two possible pathways : ( i ) Taking actions in the real world to generate new observations and ( ii ) Imagining to take actions and predicting the new observations . Consider the humanoid robot from the MuJoCo environment ( Mordatch et al. , 2015 ) . In the first case , the humanoid agent takes an action in the real environment , observes the change in its position ( and location ) , takes another step and so on . In the second case , the agent imagines taking a step , predicts what the observation would look like , imagines taking another step and so on . The first case is a close-loop setup , where the humanoid observes the state of the world , takes an action , gets the true observation from the environment , which is used to choose the next action , and so on . The second case is a open-loop setup , where the agent predicts subsequent states for multiple time steps into the future without interacting with the environment . The two cases have been summarized in figure 1 . As such , the two pathways may not to be “ consistent ” given the challenges in learning a multi-step dynamics model . By “ consistent ” , we mean the behaviour of state transitions along the two paths should be indistinguishable . Had they been consistent , the learner ’ s model would be grounded in reality ie the predictions from the open loop would be similar to the predictions from the closed loop over a long time horizon . In this work , we propose to ensure consistency by using an auxiliary loss which explicitly seeks to match the generative behaviour ( from open loop ) and the observed behaviour ( from closed loop ) as closely as possible . More generally , we show that the proposed approach helps to simultaneously train more powerful policies as well as better dynamics models , by using a training objective that is not solely focused on predicting the next observation . Our evaluation protocol consists of learning both observation-space models and state-space models . We consider various continuous control tasks from the OpenAI Gym suite ( Brockman et al. , 2016 ) , and RLLab ( Duan et al. , 2016 ) and show that using the proposed auxiliary loss consistently helps in achieving better performance across tasks . We evaluate the proposed approach on the pixel-based cheetah domain from the OpenAI Gym suite ( Brockman et al. , 2016 ) . This domain is bit difficult for the “ baseline ” state space models as one can only infer positions and not velocities from the images and hence this makes the task partially observable . We compare the proposed model to the state of the art state space models ( Buesing et al. , 2018 ) , and show that the proposed method consistently achieves better results . 2 PRELIMARIES . A finite time Markov decision processM is generally defined by the tuple ( S , A , f , R , γ ) . Here , A the action space , S is the set of states , f ( st+1|st , at ) the transition distribution , r : S × A → R is the reward function and γ the discount factor . We define the return as the discounted sum of rewards r ( st , at ) along a trajectory τ : = ( s0 , a0 , ... , sT−1 , aT−1 , sT ) , here T refers to the effective horizon of the process . The goal of reinforcement learning is to find a policy πφ that maximizes the expected return . Here φ denotes the parameters of the policy π. Model-based RL methods learn the dynamics model from the observed transitions . This is usually done with a function approximator usually parameterized as a neural network f̂θ ( st+1|st , at ) . In such case , the parameters θ of the dynamics model are optimized to maximize the log-likelihood of the state transition distribution . 3 ENVIRONMENT MODEL . Consider a learning agent which is training to optimize a reward signal r in a given environment . At a given timestep t , the agent is in some state st ∈ S. It takes an action at ∈ A according to its policy at ∼ πt ( at|st ) , receives a reward rt ( from the environment ) and transitions to a new state st+1 . The agent is trying to maximize its expected reward and has two pathways for improving its behaviour : 1 . Close-loop path : Here , the learner interacts with the environment at every step . The agent starts in state s0 and is in state st at time t. It chooses an action at to perform ( using its policy πt ) , performs the chosen action , and receives a reward rt . It then observes the environment to obtain the new state st+1 , uses this state to decide which action at+2 to perform next and so on . 2 . Open-loop path : Here , the learner predicts the future observations ( or future belief state in case of state space models ) . The agent starts in state s0 and in state st at time t. Note that the agent ” imagines ” itself to be in state sIt and can not access the true state of the environment . It chooses an action at to perform ( using its policy πt ) , performs the action in the “ learners ” world and imagines to transition to the new state sIt+1 . Thus the current “ imagined ” state is used to predict the next “ imagined ” state . During these “ imagined ” roll-outs , the agent does not interact with the environment but interacts with its “ imagined ” version of the environment which we call its dynamics model or the learner ’ s “ world ” . As an alternative , the agent could use both the pathways simultaneously . The agent could , in parallel , ( i ) Build a model of the environment and ( ii ) Engage in interaction with the environment as shown in Figure 1 . We propose to make the two pathways consistent with each other so as to ensure that the predictions from the learner ’ s dynamics model are grounded in the observations from the environment . We show that such a “ consistency constraint ” helps the agent to learn a powerful policy and a better dynamics model of the environment . 3.1 CONSISTENCY CONSTRAINT . We want the “ imagined ” behaviour ( from the open loop ) to be consistent with the observed behaviour ( from the close loop ) . Thus making sure that the predictions from the learner ’ s dynamics model ( or the “ world ” ) are similar to the actual observations from the environment . The learner ’ s dynamics model could either be in observation space ( pixel space ) or it could be in state space . State space models are generally more efficient as they model dynamics at some higher level of abstraction . In the case of state space models , the learner predicts transitions in the state space by first encoding the actual observation from the environment into the state space of the learner and then imposing the consistency constraint in the ( learned ) state space . In order to apply the consistency constraint in the open loop setup , we need to obtain the corresponding state transitions in the environment . At a given timestep t , the learner is in some environment state st while it imagines to be in state sIt . It takes an action at according to its policy at ∼ πt ( at|st ) . Now the learner can make transition in two ways . It could execute the action in the environment and transition to state st+1 ( as governed by the dynamics of the environment ) . Alternatively , it could execute the action in the “ learned ” dynamics environment f̂θ and imagine to transition to the state sIt+1 = f̂θ ( s I t , at ) . Note that the state st is not used by the learner ’ s dynamics model when making state transitions during the open-loop setup . Many possibilities exist for imposing the “ consistency constraint ” . In a simple setup , we could enforce the per step output from the open and the closed loops to be similar . The downside of this approach is that it encourages the dynamics model to mimic each and every detail of the environment irrespective of its utility in predicting the state transitions ( Lamb et al. , 2016 ) . Hence , we encode the state transitions ( during both open-loop and closed-loop ) into fixed-size real vectors using recurrent networks and enforce the output of the recurrent networks to be similar in the two cases . Encoding the sequence can be seen as abstracting out the per-step state transitions into how the dynamics of the environment evolve over time . This way , we do not focus on mimicking each state but the high level dynamics of the state transitions . This consistency loss encourages the dynamics model to only focus on information that makes the multi-step predictions ( from the open-loop ) indistinguishable from the actual future observations from the environment . This is shown in figure 1 . Once we have the predicted state transitions and the real state transitions , we could impose the “ consistency constraint ” in several ways . In order to keep our setup simple , we minimize the prediction error ( L2 error ) between the encoding of predicted future observations as coming from the learner ’ s dynamics model ( during open-loop ) and the encoding of the future observations as coming from the environment ( during closed loop ) . Let us assume that the agent started in state s0 and that a0 : T−1 denote the sequence of actions that the agent takes in the environment from time t = 0 to T −1 . Similarly , s1 : T denotes the sequence of states that the agent transitions through . Alternatively , the agent could have “ imagined ” a trajectory of state transitions by performing the actions a0 : T−1 in the learner ’ s dynamics model . This would result in the sequence of states sI1 : T . The consistency loss is computed as follows : enc ( s1 : T ) ) = RNN ( [ s1 , s2 , ... , sT ] ) enc ( sI1 : T ) ) = RNN ( [ s I 1 , s I 2 , ... , s I T ] ) lcc ( θ , φ ) = ‖enc ( s1 : T ) ) − enc ( sI1 : T ) ) ‖ ( 1 ) where ‖‖ denotes the L2 norm and RNN denotes the GRU model used to encode the sequence of state transitions into a fixed length vector . The agent which is trained with the consistency constraint is referred to as the consistent dynamics agent . The overall loss for such a learning agent can be written as follows : ltotal ( θ , φ ) = lrl ( θ , φ ) + αlcc ( θ , φ ) ( 2 ) where θ refers to the parameters of the agent ’ s transition model f̂ and φ refers to the parameters of the agent ’ s policy π . The first component of the loss function , lrl ( θ , φ ) , corresponds to the RL objective i.e maximizing expected return and is referred to as the RL loss . The second component of the loss , lcc ( θ , φ ) , corresponds to the loss associated with the consistency constraint and is referred to as consistency loss . α is a hyper-parameter to scale the consistency loss component with respect to the RL loss .

Summary: The authors propose to train a policy while concurrently learning a dynamics model. In particular, the policy is updated using both the RL loss (rewards from the environment) and the "consistency constraint", which the authors introduce. This consistency constraint is a supervised learning signal, which compares trajectories in the environment with trajectories in the imagined world (produced with the dynamics model).