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: This paper proposes the method "robust predictable control (RPC)". It proposes to use a information bottleneck between inputs and actions. However, in contrast to previous works, it does so not on a per-observation basis, but on a per-sequence basis using a latent model. Similar to some of the prior work, the encoding cost is not added to the loss, but subtracted from the reward.