Robust Predictable Control

1 Introduction . Many areas of reinforcement learning ( RL ) research focus on specialized problems , such as learning invariant representations , improving robustness to adversarial attacks , improving generalization , or building better world models . These problems are often symptoms of a deeper underlying problem : autonomous agents use too many bits from their environment . For the purpose of decision making , most information about the world is irrelevant . For example , a lane keeping feature on a car may take as input high-resolution camera input ( millions of bits ) , but only needs to extract a few bits of information about the relative orientation of the car in the lane . Agents that rely on more bits of information run the risk of overfitting to the training task . Agents that use few bits of information enjoy a number of appealing properties . These agents can better cope with high-dimensional sensory inputs ( e.g. , dozens of cameras on a self-driving car ) and will be forced to learn representations that are more broadly applicable . Agents that throw away most information will be agnostic to idiosyncrasies in observations , providing robustness to missing or corrupted observations and allowing for transfer to different scenarios . For example , if an agent ignores 99.9 % of bits , then corrupting a random bit is unlikely to change the agent ’ s behavior . Moreover , an agent that minimizes bits will prefer states where the dynamics are easy to predict , meaning that the agent ’ s resulting behavior will be easier to model . Thus , compression not only changes an agent ’ s representation , but also changes its behavior : an agent that can only use a limited number of bits will avoid risky behaviors that require more bits to execute ( see Fig . 3a ) . The generalization and robustness of a machine learning model is directly related to the complexity of that model . Indeed , standard techniques for reducing complexity , such as the information bottleneck [ 1 , 43 ] , can be directly applied to the RL setting [ 14 , 21 , 31 , 42 ] . While these approaches make 1Project site with videos and code : https : //ben-eysenbach.github.io/rpc 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . the policy ’ s action a simple function of the state , they ignore the temporal dimension of decision making . Instead , we will focus on learning policies whose temporally-extended behavior is simple , where we use compression to measure simplicity . Our key observation is that a policy ’ s behavior is simple if it is predictable . Our method improves upon prior methods that apply an information bottleneck to RL [ 14 , 21 , 31 ] by recognizing two important properties of the decision making setting . First , because agents make a sequence of decisions , they can use salient information at one time step to predict salient information at the next time step . These predictions can decrease the amount of information that the agent needs to sense from the environment . We will show that learning a predictive model is not an ad-hoc heuristic , but rather a direct consequence of minimizing information using bits-back coding [ 12 , 19 ] . Second , unlike supervised learning , the agent can change the distribution over states , choosing behaviors that visit states that are easier to compress . For example , imagine driving on a crowded road . Aggressively passing and tailgating other cars may result in reaching the destination faster , but requires careful attention to other vehicles and fast reactions . In contrast , a policy optimized for using few bits would not pass other cars and would leave a larger following distance ( see Fig . 3b ) . Combined , these two capabilities result in a method that jointly trains a latent space model and a control policy , with the policy being rewarded for visiting states where that model is accurate . Unlike typical model-based methods , our method explicitly optimizes for the accuracy of open-loop planning and results in a model and policy that are self-consistent . The main contribution of this paper is an RL algorithm , robust predictable control ( RPC ) , for learning policies that use few bits of information . We will refer to such policies as compressed policies . RPC brings together ideas from information bottlenecks , model-based RL , and bits-back coding into a simple and theoretically-justified algorithm . RPC attains up to 5× higher return than a standard information bottleneck when compared at the same bitrate . Experiments demonstrate that the compressed policies learned by RPC are more robust than those learned by alternative approaches , generalize well to new tasks , and learn behaviors that can be composed for hierarchical RL . 2 Related Work . Our work builds on a large body of prior work , both in supervised learning and reinforcement learning , that argues that compression is a good objective for learning representations and improving robustness [ 4 , 19 , 27 , 30 , 38 , 39 , 42 , 44 ] . We measure compression by the mutual information between its inputs and outputs [ 4 , 43 ] , a metric that has been used to study the generalization properties [ 43 ] and representations learned by neural networks [ 1 ] . Our work extends these results to the RL setting by observing that the agent can change its behavior ( i.e. , the data distribution ) to be more easily compressed . In the RL community , prior work has used the variational information bottleneck ( VIB ) [ 2 ] to minimize communication between agents in a multi-agent setting [ 45 ] and to improve exploration in a goal-reaching setting [ 14 ] . The most related RL methods are those that use an information bottleneck in RL to improve generalization [ 21 , 31 ] . Whereas these prior methods compress observations individually , we will compress sequences of observations , similar to a sequential VAE [ 8 , 23 ] . This difference , which corresponds to learning a latent-space model , improves compression and increases robustness on downstream tasks . Prior work in RL has sought robustness , transfer , and generalization in many ways . For example , prior RL methods often use explicit representation learning objectives to accelerate learning of image-based tasks [ 13 , 26 , 28 , 36 ] . The presentations learned by our method , like those learned by prior contrastive learning methods [ 26 , 34 , 36 ] , do not require reconstructing observations . Prior work has also studied numerous strategies for learning policies that are resilient to perturbations in the environment , a problem known as robust RL [ 20 , 22 , 33 , 41 ] . While prior robust RL methods typically involve solving a two-player game , we show that compression is a simpler mechanism for achieving some robustness benefits . Finally , the problem of learning RL policies that generalize has been studied by many prior papers [ 9 , 10 ] . 3 Reinforcement Learning with Fewer Bits . This section introduces the idea that predicting the future allows RL policies to operate with fewer bits . We derive this idea from first principles , develop it into a complete RL method , then discuss connections with model-based RL , bits-back coding , and other related topics . 3.1 Notation and Preliminaries . An agent interacts in an MDP defined by states st and actions at ; we will use “ states ” and “ observations ” interchangeably . The agent samples actions from a policy πθ ( at | st , st−1 , st−2 , · · · ) .2 We will construct this policy by learning an encoder φ ( zt | st ) ( which encodes a state st into a representation zt ) , and a high-level policy πzθ ( at | zt ) , which decodes a representation zt into an action at . The environment dynamics are defined by an initial state s1 ∼ p1 ( s1 ) and a transition function p ( st+1 | st , at ) . The RL objective is to maximize the expected γ-discounted sum of rewards : maxπ Eπ [ ∑∞ t=1 γ tr ( st , at ) ] . The discount factor can be viewed as saying that the episode terminates at every time step with probability ( 1− γ ) , an interpretation we will use in Sec . 3.2 . A model is simpler if it expresses a simpler input-output relationship [ 4 , 43 ] . We will measure the complexity of a function using an information bottleneck , which is the mutual information I ( x ; y ) between an input x and an output y [ 1 , 43 ] . The variational information bottleneck provides a tractable upper bound on mutual information [ 1 , 2 ] : I ( x ; y ) ≤ Ep ( x , y ) [ log ( p ( y | x ) m ( y ) ) ] , where m ( y ) is an arbitrary prior distribution . Applying the information bottleneck to an intermediate layer z = φ ( x ) is sufficient for bounding the mutual information between input x and output y . Following prior work on compression in RL [ 21 , 31 ] , we aim to maximize rewards while minimizing the number of bits . While there are many ways to apply compression to RL ( e.g. , compressing actions , goals , or individual observations ) , we will focus on compressing sequences of states . The input is a sequence of states , s1 : ∞ , ( s1 , s2 , · · · ) ; the output is a sequence of actions , a1 : ∞ , ( a1 , a2 , · · · ) . The objective is to learn a representation φθ ( zt | st ) and policy πθ ( at | zt ) that maximize reward , subject to the constraint that the policy uses on average C > 0 bits of information per episode : max θ Eπ , φ [ ∞∑ t=1 γtr ( st , at ) ] s.t . Eπ [ I ( s1 : ∞ ; z1 : ∞ ) ] ≤ C , ( 1 ) While prior work on compression in RL [ 21 , 31 ] has applied the VIB to states independently , we will aim to compress entire sequences of observations . Applying the VIB to sequences allows us to use more expressive choices of the prior m ( z1 : ∞ ) . Our prior will use previous representations to predict the prior for the next representation , allowing us to obtain a tighter bound on mutual information . 3.2 Using Fewer Bits by Predicting the Future . Predicting the future allows agents to observe fewer bits from future observations . To further decrease the number of bits used , the agent can change its behavior to visit states that are more easily compressed . States where the dynamics are hard to predict will require more bits , so the agent will prefer visiting states where its learned model can accurately predict the next state . This intuition corresponds to solving the optimization problem in Eq . 1 with a prior that is factored autoregressively : m1 : ∞ ( z1 : ∞ | a1 : ∞ ) = m ( z1 ) ∏ tmθ ( zt+1 | zt , at ) . Note that the prior has learnable parameters θ . We apply the VIB to obtain an upper bound on the constraint in Eq . 1 : Eπ [ I ( s1 : ∞ ; z1 : ∞ ) ] ≤ Eπ [ log ( p ( z1 : ∞ | s1 : ∞ ) m ( z1 : ∞ ) ) ] = Eπ [ ∞∑ t=1 γt ( log φθ ( zt | st ) − logmθ ( zt+1 | zt , at ) ) ] . ( 2 ) This objective is different from prior work that applies the VIB to RL [ 14 , 21 , 31 ] because the prior mθ ( zt ) is predicted , rather than fixed to be a unit Normal distribution . The discount factor γ reflects the assumption from Sec . 3.1 that the episode terminates with probability ( 1− γ ) at each time step ; of course , no bits are used after the episode terminates . Our final objective optimizes the policy π , the encoder φ , and the prior m to maximize reward and minimizing information : max θ Eπθ , φθ , mθ [ ∞∑ t=1 γtr ( st , at ) ] s.t . Eπθ , φθ , mθ [ ∑ t γt ( log φθ ( zt | st ) − logmθ ( zt+1 | zt , at ) ) ] ≤ C. ( 3 ) For the encoder , this objective looks like a modification of the information bottleneck where the prior is predicted from the previous representation and action . Of course , we can not predict the first 2We condition the policy on sequences of states because , even if the task is Markovian , the agent ’ s choice to omit certain bits from the representation can turn an otherwise fully-observed task into a partially-observed task . representation , so we substitute use a unit Normal distribution for the prior at initial time step . For the policy , this objective looks like the standard RL objective but with a different reward function , one that includes an additional information cost : r̃λ ( st , at , st+1 ) , r ( st , at ) + λ ( logmθ ( zt+1 | zt , at ) − log φθ ( zt+1 | st+1 ) ) ︸ ︷︷ ︸ ( negative of the ) information cost , ( 4 ) where λ is the cost per bit . The term log φθ ( zt | st ) corresponds to the number of bits required to represent the representation zt . The term logmθ ( zt+1 | zt , at ) reflects how well the agent can predict the next representation . In effect , the agent has to “ pay ” for bits of the observation , but it gets a “ refund ” on bits that it predicted from the previous time step , analogous to bits-back coding [ 12 , 19 ] . Said in other words , we want to minimize the excess bits required to infer zt from the current state , relative to the number of bits required to predict zt from the previous ( zt−1 , at−1 ) . Note that the standard VIB approaches to RL [ 14 , 21 , 31 ] do not receive this refund because they compress observations independently , rather than compressing sequences of observations . Note that the policy , encoder , and model are all jointly optimized with respect to the same objective , so updates to any component increases the objective for all components . Importantly , the agent optimizes not only its representation but also its behavior to minimize this information cost : the agent learns a representation that is easily predictable and learns to visit states where that representation is easily predictable . We therefore call our method robust predictable control ( RPC ) . 4 A Practical Algorithm Our method for optimizing ( Eq . 3 ) is an actor-critic method applied to the informationaugmented reward function in Eq . 4 . We introduce a Q function , Qψ ( st , at ) = E [ ∑ t γtr̃ ( st , at , st+1 ) ] , corresponding to the augmented reward function r̃ . We optimize this Q function using standard temporal difference learning : L ( ψ ) = 1 2 ( Qψ ( st , at ) − yt ) 2 ( 5 ) where yt = br̃t + γQψ ( st+1 , at+1 ) csg and b·csg is the stop-gradient operator [ 6 , 11 ] . Since our overall objective ( Eq . 1 ) only entails compressing the policy , not the Q function , we condition the Q function directly on the state . To optimize the encoder , prior , and policy , we reexpress our objective ( Eq . 1 ) in terms of the immediate reward plus the Q function ( see derivation in Appendix A.1 ) : L ( θ ; st ) = Ezt−1∼φθ ( st−1 ) , st∼p ( st|st−1 , at−1 ) zt∼φ ( st ) , at∼πzθ ( at|zt ) [ Qψ ( st , at ) +λ ( logmθ ( zt | zt−1 , at−1 ) − log φθ ( zt | st ) ) ] . ( 6 ) Since the encoder is stochastic , we compute gradients using the reparametrization trick [ 25 ] . The fact that all three components are jointly optimized with respect to the same objective makes implementation of RPC surprisingly simple . Unlike sequential VAEs and RNNs , training agents with RPC does not require sampling frommθ ( zt+1 | zt , at ) and does not require backpropagation through time . In our implementation , we instantiate the encoder φθ ( st ) , the prior mθ ( zt+1 | zt , at ) , and the high-level policy πzθ ( at | zt ) as neural networks with parameters θ . The Q function Qψ ( st , at ) is likewise represented as a neural network with parameters φ . We update the dual parameter λ ≥ 0 using dual gradient descent . We use standard tricks such as target networks and taking the minimum over two Q values . We refer the reader to Appendix B and the open-sourced code for details .

Summary: In this paper, the authors propose a new reinforcement learning method, called robust predictable control, that optimizes a policy under limited representational resources. Similar to previous methods, they formalize the problem in a fully observable MDP and learn a compressed representation of the state space using a variational information bottleneck constraint. Crucially, they apply the bottleneck on sequences of states. Through a clever factorization of the variational prior, they obtain a prior dynamics model in the compressed state space. This dynamics model is predictive and the agent only needs additional representational resources if the encoded states are not predicted by the prior dynamics already. Using a set of standard RL benchmarks, the authors show that this set up improves compression and enforces predictable and robust behaviors .