Title: Robust Predictable Control
Authors: Benjamin Eysenbach
Ruslan Salakhutdinov
Sergey Levine
Venue: NeurIPS 2021
Keywords: reinforcement learning
information bottleneck

Abstract: Many of the challenges facing today's reinforcement learning (RL) algorithms, such as robustness, generalization, transfer, and computational efficiency are closely related to compression.  Prior work has convincingly argued why minimizing information is useful in the supervised learning setting, but standard RL algorithms lack an explicit mechanism for compression. The RL setting is unique because (1) its sequential nature allows an agent to use past information to avoid looking at future observations and (2) the agent can optimize its behavior to prefer states where decision making requires few bits. We take advantage of these properties to propose a method (RPC) for learning simple policies. This method brings together ideas from information bottlenecks, model-based RL, and bits-back coding into a simple and theoretically-justified algorithm. Our method jointly optimizes a latent-space model and policy to be self-consistent, such that the policy avoids states where the model is inaccurate. We demonstrate that our method achieves much tighter compression than prior methods, achieving up to 5$\times$ higher reward than a standard information bottleneck when constrained to use just 0.3 bits per observation. We also demonstrate that our method learns policies that are more robust and generalize better to new tasks.

TL;DR: We propose a method for learning robust and predictable policies in RL using ideas from compression.

Review discussion:
**Originality**
* The paper clearly states the difference to prior work and adequately cites it.
* In my opinion, the paper has close connection to Rate-Distortion-Theory, or,
to be more precise, "bounded rational decision-making".
Especially the experiment in Figure 3a shows this behaviour quite nicely.
I think the paper would benefit from establishing the connection and referring to this domain of research.
**Quality**
* I perceive the paper as of very high quality.
It is a consistent paper containing a sound theoretical motivation of the
followed approach, interesting and well-executed theoretical derivations
and insightful experiments.
Prior work is adequately cited and compared to in the experimental evaluation.
**Clarity**
* In general, the paper is very clearly written and good to follow. I have the following comments:
* The term "states" and "observations" seem to be used interchangeably.
I would appreciate a clarification of that fact at the beginning of section 3.1, eventually
including the statement that you assume the environment state to be fully observable, if this is the case.
* Related to the first point, in l. 95 you define the policy to be conditioned on
a *sequence* of past states. If $s_t$ are *states* of an MDP (and not observations of a POMDP)
then conditioning on the current state should be sufficient.
* The joint distribution $m_{1:\infty}(z_{1:\infty})$ as defined in l. 127 misses a factor to be a proper joint distribution
in my opinion, e.g. $m_{1:\infty}(z_{1:\infty}) = p(z_1) \prod_{t \geq 1} m(z_{t+1} | z_t, a_t)$
* The definition of the probabilistic bound in l. 242 is not aligned with the introducing text.
Due to the absolute difference it is rather the "probability that the policy's expected
return on an observed (deterministic) MDP is much *different* to the policy's expected return ..."
* Appendix, l. 694/695: Introducing $p$ and $q$ here as done in l. 710 would definitely help
to understand the following derivations.
Furthermore, I think that l. 695 contains a typo and it should
be $\log p - \log q$ in the expectation.
* Appendix, l. 698-699: I think the statement should be slightly changed to clarify that
"expectations over a mixture of finite-horizon episodes are equivalent
to expectations over a discounted infinite-length episode"
* Minor comments:
* l. 133: Typo, you probably mean "prior $m$"?
* l. 166: Typo, $z_{t-1} \sim \Phi(s_{t-1})$
* eq. 4 / l.139: time-index in prior $m$ not consistent
* l. 241: Introduce the "observed MDP" as $\hat{\mathcal{M}}$
**Significance**
In my opinion, the paper is of very high significance for the RL community as it contains insightful theoretical and experimental results on applying ideas of compression / information-regularization to reinforcement learning.
It establishes interesting connections between open-loop and closed-loop control and model-free and model-based RL.
**Conclusion**:
To me, this paper is a clear accept. It is well-written and contains a well-motivated, novel approach with interesting theoretical and experimental results for the reinforcement learning community.
No
Thank you for the detailed and insightful review.
> connection with bounded rationality
Thanks for this suggestion. We will include a discussion of this in the final paper. We welcome any suggestions of particular prior work on bounded rationality to discuss.
> states vs observations
We will clarify this as suggested.
> policy conditioned on a sequence of past states
We introduce the policy as being conditioned on 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. We will clarify this in the main text.
> l. 127 misses a factor to be a proper joint distribution
> l. 242 is not aligned with the introducing text
> I think that l. 695 contains a typo and it should be log ⁡p − log ⁡q in the expectation.
> Minor comments: ...
Thanks for pointing out these typos. We will fix them in the final version of the paper.
> Introducing p and q here as done in l. 710 would definitely help to understand the following derivations.
> Appendix, l. 698-699: I think the statement should be slightly changed to clarify that ...
Thanks for these suggestions. We will do this.
Overall, I think this is an interesting and valuable paper which is well written and easy to read.
While the idea itself does not seem like
a too big departure from previous work (although it is novel to my knowledge), the paper provides a
lot of analysis (both experimental and theoretical) as well as an intuitive explanation of how the
newly proposed reward affects the behavior of the agent - which I believe is a valuable contribution.
While I'm in favor of acceptance, I'm hesitant to endorse it stronger at the moment as it does not
provide clear information about the absolute achievable return for some of the experiments (at least
as far as I can tell).  And while it is expected that the provided information-regularisation
decreases the training performance, the extend to which this happens is important to judge the
validity of the proposed method.
In particular, it would be important to know whether the shown
advantages of RPC hold even with (somewhat) competitive performance on the task or if they are only
realised in the very low information & return regime.
To some extend the information is provided (e.g. Fig 6/8) but those plots are cut off and don't show
the performance of the unperturbed agent and also seem to use some different bitrates, all on the
very low end of the spectrum.
In short: It would be helpful to have a overview table how RPC
compares to SAC for different bitrates on all used environments.
On a positive note, I think section 5 is great. I did have a look at the proofs and
they seem reasonable, although it's possible I overlooked something.
Minor comments:
* Fig 3b-top) Why is the variance of the blue line so much higher?
* Appenidx line 695: This should be a `- log q(...)` not `- log p(...)`?
* Eq. 10: Should be `phi(z_{t+1} | s_{t+1})` ?
No
Thanks for the detailed and thoughtful review.
> It would be helpful to have a overview table how RPC compares to SAC for different bitrates on all used environments.
The bitrate constraint is slack when we run RPC with a bitrate constraint of 3 bits.
Thus, SAC should perform identically to RPC with a bitrate of 3 (see, e.g., Fig 10 in the Appendix). We will include a plot comparing with SAC in the final paper.
> Fig 3b-top) Why is the variance of the blue line so much higher?
The variance here indicates that the distance between the learned agent and the lead car changes throughout an episode. This can happen, for example, if the agent in front suddenly breaks.
> Typos
Thank you for pointing out these typos. We will fix them in the final version of the paper.
Thank you for your clarifications!
All agents face a fundamental problem: how much effort should they invest in order to achieve a certain outcome? In RL, the quality of an outcome can be described by the collected rewards, whereas the effort can be described by the informational expenses. The authors tackle this problem by considering an important informational expense, namely the complexity of the representation on which an agent bases its decisions. Low complexity representations are desirable because they are cheap (e.g.
in terms of storage) but also because they tend to avoid overfitting. If agents don’t overfit they are more robust and can generalize better. The authors thus address a fundamental issue of decision-making and show that their method actually achieves improvements in compression, robustness, and generalization.
While previous work has addressed the same problem in a similar setups (e.g. ref 25), the authors here propose a valuable extension.
Instead of compressing states one-by-one, they compress state sequences as a whole and thus get a step closer to the actual problem faced by the agent.
Furthermore, the predictive prior this approach yields is useful on its own as something akin to a default dynamics model and the authors showcase one application of this in hierarchical RL (fig. 5). Most importantly, their approach also enforces the agent to learn behaviors that are inherently predictive. I think this is a conceptually important point and they nicely show evidence for this in fig. 3.
The only issues I have are with clarity of exposition:
- Eq. 2: It is unclear whether there is really equality here. In my understanding every sequence, or partially truncated sequence, has a certain mutual information associated.
When taking this MI as a cost, we can sum them up similar to rewards and take into account the discount factor in the sense of a probability of a sequence continuing.
When doing this we get the RHS, but this is not equal to the LHS.
Furthermore, the prior that is factored autoregressively should read $m(z_{1:\inty}|a_{1:\infty})$, i.e. $m(z_{1:\inty}|a_{1:\infty}) = \prod_t m(z_{t+1}|z_t,a_t)$, or otherwise one would have to sum over probability of actions.
- Eq. 5: It is unclear where this objective is coming from and how it is optimized (I also couldn’t find the code on the provided website to check). The Q function is learned online by TD learning minimizing $L(\psi)$. Since $L(\theta)$ is not considering whole sequences, is it really a function depending on $s_{t-1}$, i.e. $L(\theta,s_{t-1})$? And then this function is optimized for all $s_{t-1}$? It is unclear to me how this optimization procedure works, and how does it relate to a policy gradient approach as in ref. 25 or as in actor critic methods?
- Some sloppy notation or typos. E.g. on bottom of page 4, the high-level policy is denoted $\pi(a|s)$, but shouldn’t it be $\pi(a|z)$ as in fig. 1?
Similarly, in the variables over which the expectation is taken in eq. 5 it reads $z_{t-1} \sim \phi(z_{t-1})$ instead of $z_{t-1} \sim \phi(s_{t-1})$. On page 3 the prior is suddenly called $r$. In eq. 4 the regularized reward function should depend on $z_t, z_{t-1},a_{t-1}$ as well. These are just some I found.
- It would be interesting to see some visualizations of the representations and dynamics that RPC finds. E.g. for the hierarchical tasks, $z$-space is three dimensional and could be visualized I assume. Visualizations would give a better intuition for how RPC exactly works.
No
Thank you for the detailed review and insightful comments.
> When doing this we get the RHS, but this is not equal to the LHS
We didn't follow this point. Would it be possible to elaborate on it?
> the prior that is factored autoregressively should read ...
Yes, thank you for catching this. We will update the paper to fix this.
> It is unclear where this objective (Eq. 5) is coming from.
Our aim is to maximizing the sum of (information-augmented) rewards (Eq. 3), starting at state $s_t$:
\mathcal{L}(\theta; s_t) = E\left[ \sum_t \tilde{r} \mid s_t \right]
= E\left[ \sum_{t'=t}^\infty \gamma^{t' - t} (r(s_{t'}, a_{t'}) + \lambda ( \log m_\theta(z_{t'+1} \mid z_{t'}, a_{t'}) - \log \phi_\theta(z_{t'} \mid s_{t'}))) \mid s_t \right] $$
$$= E\left[ r(s_t, a_t) + \lambda ( \log m_\theta(z_{t+1} \mid z_t, a_t) - \log \phi_\theta(z_t \mid s_t)) + \sum_{t'=t+1}^\infty \gamma^{t' - t} (r(s_{t'}, a_{t'}) + \lambda ( \log m_\theta(z_{t'+1} \mid z_{t'}, a_{t'}) - \log \phi_\theta(z_{t'} \mid s_{t'}))) \mid s_t \right] $$
$$= E\left[ r(s_t, a_t) + \lambda ( \log m_\theta(z_{t+1} \mid z_t, a_t) - \log \phi_\theta(z_t \mid s_t)) + \gamma Q_\psi(s_{t+1}, a_{t+1}) \right].$$
Note that the reward at the current time step, $r(s_t, a_t)$, is not influenced by the parameters $\theta$, so we can drop this term:
$$\stackrel{const.}{=} E\left[\lambda ( \log m_\theta(z_{t+1} \mid z_t, a_t) - \log \phi_\theta(z_t \mid s_t)) + \gamma Q_\psi(s_{t+1}, a_{t+1}) \right].$$
The remaining difference between this objective and Eq. 5 is that the Q value term is scaled by $\gamma$.
However, this difference has no effect on the optimization problem because the parameter $\lambda$ is automatically tuned. If we scale the second term, $Q$, by some value (say, $\gamma$), then tuning $\lambda$ to satisfy the bitrate constraint will result in a different value for $\lambda$ (one which is $\gamma$ times smaller).
We will update the text (1) to include this derivation, (2) to note that $\mathcal{L}(\theta; s_t)$ should also depend on $s_t$, (3) and to clarify how the states $s_t$ are sampled (from the replay buffer).
> how does it relate to a policy gradient approach
The gradient of this objective w.r.t. the policy is almost identical to the standard policy gradient (see, e.g., [Sutton and Barto, Page 334 - 335]). The difference is that the Q function used in our paper reflects the information-augmented reward $\tilde{r}$, rather than the task reward.
> Typos
Thank you for pointing out typos. We will fix them in the final version.
> visualizations of the representations and dynamics that RPC finds
We included some visualizations of the learned representations in Fig. 4. We will create videos of the dynamics of the learned representations and include them in the final paper.
I am generally accepting of the work. I would appreciate discussing the following issues.
The authors admit that being compressive will necessarily lead to lower rewards in complex situations.
But then it felt like you took it as a goal because it by itself is still a good thing to achieve.
Later you show that they have better generalization power.
But can you motivate a bit more why aiming for compressed policies is a good goal? You gave some specific examples where compressed policies will naturally be safer in some cases.
But if we are trying to build a general learning agent, where the agent needs to learn many things and address more and more complex situations over time, should we aim for compressed policies?
In equation 3, we have phi(z|s) and m(z’|z,a) in pairs.
However, in the rest of the paper, phi(z’|s’) and m(z’|z,a) come in pairs. Seems like there is a mismatch. Can you explain the mismatch?
Clarification questions:
A large part of the paper initially talks about learning a “simple policy,” which is defined as a simple input-output relationship. It will be good if you specify directly that by simple you mean compressive.
You claimed that “Q value needs at least as many bits as a representation sufficient for predicting the action.” This argument may make sense for value-based methods such as Q-learning where the policy is derived from learned values.
But in your case, the policy is parameterized directly and not through the values function. Can you clarify why this argument still makes sense for parameterized policies?
When you compare with VIB variants, did you make sure that they use the same amount of computation and memory (e.g. network size) as RPC? If not, how did you make sure that this comparison is fair?
In Figure 5, the x-axis of the plots has episodes. Are they of the same length? If not, how is it a sensible choice to align in terms of episodes?
And if those tasks have constant episode length, then it is a fixed horizon problem. Did you pass the remaining time steps as part of the observation?
Otherwise, is the task from the agent’s perspective really well defined that we can describe coherently?
Why action repeat is a reasonable or interesting baseline?
Few edit suggestions:
Line 32: “minimizing bits” -> “minimizing the number of bits”
Line 157: both Q and y are functions of psi. This means both will be differentiated. Is that what you did because that’s not what temporal difference learning does.
No
Thank you for the very detailed review.
> can you motivate a bit more why aiming for compressed policies is a good goal?
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 (e.g., [1, 2, 3, 4, 5, 6, 7, 8]).
For example, [4] argue that "Supervised neural networks generalize well if there is much less information". [6] similarly argue that "compression provides a natural quantitative approach to the question of relevant information." [8] notes that "‘simpler’ models should lead to better generalization." We will revise the discussion of compression to highlight this long history of prior work.
> If we are trying to build a general learning agent, where the agent needs to learn many things and address more and more complex situations over time, should we aim for compressed policies?
While we do not claim to build a general learning agent, our hierarchical experiments and similar experiments from prior work on compression in RL [3, 9] suggest that compression might be an effective tool for adapting to new and more complex tasks over time.
> $\phi(z \mid s)$ vs $\phi(z' \mid s')$
We apply the same encoder $\phi(z \mid s)$ to every state: $z \sim \phi(z \mid s)$, and $z' \sim \phi(z' \mid s')$. We will update the paper to make this consistent.
> specify directly that by simple you mean compressive.
We will clarify this.
> conditioning Q values on the representation
As noted by the reviewer, our objective (Eq. 3) only entails compressing the policy, not the Q function. The point we attempted to make in L163 is that we do *not* need to compress the Q  function. Learning a representation that is sufficient for both the policy and the Q function would lead to worse compression.
> did you make sure that they use the same amount of computation and memory (e.g. network size) as RPC?
We used exactly the same hyperparameters and architectures for RPC and the VIB baselines. The policies (and Q functions) used by both methods had exactly the same number of parameters.  Please see Appendix B for details.
> episodes length in Fig 5
Yes, the episode lengths are fixed for each environment (see Appendix B.7). The same data plotted against environment steps would therefore look identical.
> fixed horizon problem
We optimize over a *sequence* of representations, $[z_1, z_2, \cdots]$.
That is, each representation is associated with a particular index in the sequence. The optimizer therefore (implicitly) knows the time at which each behavior will be executed.
> Why is action repeat a reasonable or interesting baseline?
In some tasks, actions correspond to a desired pose.
Thus, action repeat would correspond to moving to a desired pose, which seemed like a reasonable baseline in these manipulation tasks. Our experiments highlight that the behaviors learned do more than move to a particular pose.
In addition, they do obstacle avoidance and objective manipulation.
> Edits
Thanks for the suggestions! We'll incorporate them in the final version.
> Line 157: both Q and y are functions of psi…
Yes, there should be a "stop gradient" (a.k.a. detach()$ around the target Q values. We will update the paper.
[1] Veness, Joel, et al. "Compress and control." Twenty-Ninth AAAI Conference on Artificial Intelligence. 2015.
[2] Bassily, Raef, et al. "Learners that use little information." Algorithmic Learning Theory. PMLR, 2018.
[3] Schmidhuber, Jürgen. "Learning complex, extended sequences using the principle of history compression." Neural Computation 4.2 (1992): 234-242.
[4] Hinton, Geoffrey E., and Drew Van Camp. "Keeping the neural networks simple by minimizing the description length of the weights." Proceedings of the sixth annual conference on Computational learning theory. 1993.
[5] Fischer, Ian. "The conditional entropy bottleneck." Entropy 22.9 (2020): 999.
[6] Tishby, Naftali, Fernando C. Pereira, and William Bialek. "The information bottleneck method." arXiv preprint physics/0004057 (2000).
[7] Littlestone, Nick, and Manfred Warmuth. "Relating data compression and learnability." (1986).
[8] Shamir, Ohad, Sivan Sabato, and Naftali Tishby. "Learning and generalization with the information bottleneck." Theoretical Computer Science 411.29-30 (2010): 2696-2711.
[9] Schwarz, Jonathan, et al. "Progress & compress: A scalable framework for continual learning." International Conference on Machine Learning. PMLR, 2018.
It will be appreciated if you also provide some intuition on why compression can be good in the context of learning policies.
It is more than $\phi(z|s) $ vs $\phi(z'|s')$. I also pointed out that In equation 3, you have $\phi(z|s)$ and $m(z’|z,a)$ in pairs.
However, in the rest of the paper, $phi(z’|s’)$ and $m(z’|z,a)$ come in pairs. If I try to account for all the $\phi$s and $m$s, then it appears that the first $\phi(z|s)$ is missing.
Pseudocode will make the algorithm much more easily understandable.
"Yes, there should be a "stop gradient" (a.k.a. detach()$ around the target Q values. We will update the paper."
That will be much appreciated. Note that `detach` isn't a mathematical operation. It will be helpful to find a way of expressing it mathematically.
Thanks for the quick reply!
> It will be appreciated if you also provide some intuition on why compression can be good in the context of learning policies.
When starting this project, our hypothesis was that compression would be good in the context of learning policies because it would result in "simpler" policies, and that these simpler policies would be more robust, more predictable, and more easily composed for hierarchical RL.
For example, we expected compressed policies to be more robust because they look at only a few bits of the observation, so most perturbations to the observation would not affect the policy's actions. Our experiments suggest that this hypothesis is true.
> $\phi(z \mid s)$
Thank you for clarifying the feedback here. The mismatch is caused because we evaluate the encoder *one more time* than we evaluate the model.
For example, in an episode where the policy visits $H$ *states* we evaluate the encoder on each of these $H$ states but only evaluate the model $H - 1$ times. The root cause is that we cannot use the model to predict the first representation, $z_1$.
Instead, we use a fixed prior $N(0, 1)$ for the representation at the first time step. We will clarify this in the main text.
> Pseudocode will make the algorithm much more easily understandable.
Thank you for the suggestion. We will include a pseudocode Algorithm box in the final version.
> Note that detach isn't a mathematical operation. It will be helpful to find a way of expressing it mathematically
We welcome suggestions for explaining this more clearly (e.g., [1, 2] use $\perp$ for this symbol, [3, 4] use $STOPGRAD(\cdots)$, [5] use $|\cdot|_{\times}$, [6] use $detach(\cdots)$).
[1] Foerster, Jakob, et al. "Dice: The infinitely differentiable monte carlo estimator." International Conference on Machine Learning. PMLR, 2018.
[2] Xu, Zhongwen, et al. "Meta-gradient reinforcement learning with an objective discovered online." arXiv preprint arXiv:2007.08433 (2020).
[3] Chen, Xinlei, and Kaiming He. "Exploring simple siamese representation learning." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021.
[4] Arpit, Devansh, et al. "h-detach: Modifying the LSTM gradient towards better optimization." arXiv preprint arXiv:1810.03023 (2018).
[5] Fujimoto, Scott, David Meger, and Doina Precup. "An equivalence between loss functions and non-uniform sampling in experience replay." arXiv preprint arXiv:2007.06049 (2020).
[6] Shekhovtsov, Alexander, Viktor Yanush, and Boris Flach. "Path Sample-Analytic Gradient Estimators for Stochastic Binary Networks." arXiv preprint arXiv:2006.03143 (2020).
Thanks for your detailed explanations.
Regarding `detach`, as long as the tools we are using to explain this concept are defined mathematically, it should be fine.
> Regarding detach, as long as the tools we are using to explain this concept are defined mathematically, it should be fine.
We will make sure to do that.

Meta-review:
role: PC
data: comment: All 4 reviewers suggested acceptance of the paper and the authors clarified all open questions and concerns in their rebuttal. Therefore I am recommending acceptance of the paper.

Decision: Accept (Spotlight)