Maximum Total Correlation Reinforcement Learning

Thank you for thoroughly reviewing our work and appreciating the performance and robustness of our method, and the persuasive evidence of our main claims.

benefit of total correlation maximization in non-periodic tasks

Our main hypothesis that simple behavior that does not overfit to slight variations is less brittle is not restricted to periodic tasks, but also well-motivated for non-periodic tasks, such as manipulation tasks. For example, consider two different policies for grasping an object. The first policy is able to achieve high reward by always using essentially the same grasping motion and no adaptation. The second policy achieves similar reward, but relies more strongly on feedback and adapts even to slight variations. We expect the first behavior, which would have higher total correlation and better predictability but no periodicity, to be more robust to perturbations that have not been observed during training.

more non-periodic experiments

We conducted additional experiments on five non-periodic MetaWorld tasks. On all tasks we achieved similar or better success rate than the baselines, supporting our hypothesis that regularizing towards simpler behavior is also beneficial for non-periodic tasks. A table with the results can be found in our reply to reviewer ts9N.

Derivation of eq (11)

Our derivations do not assume that $\log(p(z\_t))$ is equal to $\log(p(z\_t|s\_t))$. However, due to the non-negativity of the expected KL $E\_{p(s\_t)} [ \text{KL}(p(z\_t|s\_t) || p(z_t))]$, substituting $p(z\_t|s\_t)$ for $p(z\_t)$ or $\pi(a\_t|s\_t)$ for $p(a\_t)$, does not invalidate our lower bound.

magnitude of noise

Thank you for this constructive feedback. We use stronger state and action noises now. The following tables compare the performance (mean and 90% confidence interval over 20 seeds) on all eight tasks. Overall, MTC obtains better performance in the presence of strong noise.

Inconsistency between Figure 2 and Figure 7

The inconsistency is caused by using different $I_p$, which we did not tune for each task in our robustness experiments.

depicting behavior in non-periodic environments

As we are not able to show a graph in our reply, we evaluated the action-predictability to demonstrate that our regularizer also improves simplicity for non-periodic tasks. The table can be found in our reply to Reviewer 8GSd.

Related approaches using Fourier transform [1,2]

Thank you for suggesting these methods, which are related to MTC and will be discussed in our revision. However, they are not closely related, since they do not consider the simplicity of behavior.

Normalized scores and learning curves.

We can not show the learning curves here, but will add them to the revision with original rewards. MTC shows good learning stability.

Figure 7

We will put Figure 7 in the main in our revision.

Figure 4

No, Fig. 4 does not show robustness in non-periodic tasks. We focused on DMC for our robustness experiments to follow the setup of closely related work (RPC and LZ-SAC).

Thank you for carefully reviewing our submission, and acknowledging the soundness of the main experiments, the good presentation and the empirical results that show more robust and compressible policies.

the claim of reducing spurious correlation should be investigated more directly or removed

We evaluated the effect of spurious correlation in Appendix C4. We did not add Gaussian noise to the observation, but added additional state dimensions that are not controllable by the actor, but instead follow a fixed Gaussian transition model. These distractors are not correlated with the remaining states, the actions nor the reward that the agent receives, although by coincidence, it might appear that such correlations exists, resulting in spurious correlations. Our results show that MTC achieves higher rewards than RPC and SAC on the Walker Stand task, suggesting that MTC improves robustness to spurious correlations.

Some experiments related to the evaluation of correlation could be done more directly

Thank you for suggesting evaluating the predictability to more directly assess the total correlation. We trained a t-step-ahead predictive model for t=[3,5,8,10] time steps on data sets that have been collected with the policy that has been learned with the different methods for a locomotion task and a manipulation task. The model was trained to predict the action t steps ahead, when given the current action. For all tested time differences, the prediction error of MTC was the smallest among all baselines.

non-markovianness of the reward as well as the non-stationarity of the reward

Although we clearly state that our augmented reward depends on the history-based decoder, we agree that the notation should reflect this dependency, and that the implication for the optimization should be discussed, and we will revise our submission accordingly. Our reward function depends on past states and actions to encourage the agent to act consistently with its previous actions. For our experiments, however, we did not provide the past states and actions to the policy, which limits its ability to perform optimally due to the non-Markovianity of its state-action-space. We did not provide the history to the policy solely to improve the simplicity of our method and the fairness of the comparisons. However, we performed an additional experiment on the Finger Spin task for the rebuttal to evaluate our approach in a Markovian setting (by using a recurrent policy), which did not result in significant differences.

We also agree that the non-stationarity of the reward should be discussed and we will revise the submission accordingly. Non-stationary rewards are quite common nowadays and are used effectively in different subfields, such as imitation learning, representation learning, and---most-relevant for us---by related works such as RPC and LZ-SAC. While establishing convergence guarantees for MTC (and the related works) would be desirable, we unfortunately have to leave this for future work.

Thank you for carefully reviewing our submission, and appreciating our contributions, and thorough experimentations.

is MTC better because of total correlation, or because of it being non-Markovian?

While our total correlation objective results in a reward function that depends on the past in order to encourage the agent to perform actions that are consistent with its previous actions, we did not choose a history-based policy in our experiments for a simpler and fairer comparison. We ran additional experiments to test the effect of history-based observations for standard SAC and our method on the Finger Spin task. The results are consistent with non-history-based policies, indicating that the improvements are caused by our total correlation objective, whereas non-markovianity is not sufficient.

Robustness claims would be more solid if the other algorithms did not start worse off

Maintaining high performance in the face of noise or dynamics mismatch is significantly harder than maintaining low performance. Indeed, the performance of the worst possible policy could not degrade when subjecting it to observation noise. Instead, we argue that it is more relevant to consider the absolute performance in such test settings, rather than considering the changes relative to the training performance.

The main hypothesis, a policy that exhibits more periodic behavior is less brittle, needs better substantiation.

We want to clarify that our primary interest is in learning simple behavior, which may, but does not need to manifest in periodic behavior. We challenged the hypothesis that simpler behavior tends to be less brittle on various periodic and non-periodic tasks, but did not encounter detrimental effects. We did not intend to claim that periodicity always increases robustness, and will carefully revise our submission to improve clarity. We will also revise the introduction to better motivate our hypothesis:

During training, the reinforcement learning agent, for example a robot, typically observes many slight variations in the state, for example due to sensor noise or unmodelled dynamic effects. In some cases, it can be crucial for the agent to drastically adapt its behavior to react to these variations, but sometimes they can be safely ignored. We aim to bias the agent towards ignoring such variations, whenever doing so does not seriously decrease the expected return, and thereby follow Newton's first rule of his third book of principia: "We are to admit no more causes of natural things than such as are both true and sufficient to explain their appearances". Intuitively, we expect a given behavior that performed well for previous variations, to also perform well for future variations. We introduce this inductive bias by means of the additional objective of maximizing the total correlation within the trajectory produced by the agent. This total correlation corresponds to the amount of information that we can save by using a joint encoding of all (latent) states and actions within trajectories, compared to compressing all time steps independently. By maximizing total correlation, the agent is encouraged to produce compressible and predictable trajectories, and thereby biased towards open-loop behavior such as clean periodic gaits, without preventing it from performing adaptations when necessary.

why there is the term $r(s_T, a_T)$ in equation (3)?

Thank you for spotting this mistake. The term belongs outside of the summation.

I am not so sure about the use of the word "Coherence" for trajectories that show periodic behaviour or more total correlation

Thank you for this feedback. We used coherence mostly interchangeably with consistency, but will only refer to consistency in the revision.

[why rounding is necessary]

From the point of lossless compression, the least significant bits are just as relevant as the most significant ones. When directly writing the floating point numbers of the simulator into a file, most of the bits in that file are used to encode digits that are beyond the precision of the control. When compressing such a file, the resulting file size will largely depend on how effectively those essentially random bits can be compressed, which significantly increases the variance of the resulting file size.

Is periodicity -> robustness?

We do not believe that periodicity implies robustness and did not intend to make such a claim. We hypothesize that an inductive bias towards simpler behavior that avoids unnecessary adaptions can often improve robustness.

Thank you for providing valuable comments and acknowledging the novelty and empirical performance of our approach.

Method uses a lower bound for total correlation that is always negative, making it unclear how accurately it reflects the true value.

As stated under limitations, our lower bound is not meant for approximating the total correlation, but still effective for optimization, as demonstrated in our experiments that show simpler, more compressible, and more robust behavior. To further investigate the ability of our method to maximize the total correlation, we conducted additional experiments on the Finger spin task for this rebuttal. Namely, we collected data sets using the policies that have been learned with the different methods, and trained a model to predict the action t steps ahead, based on the current action. As indicated in the following table, our lower bound object significantly increases the predictability, further suggesting that MTC is effective in increasing the consistency, predictability and thereby total correlation of the emerging behavior.

Not fully clear if the same benefits would apply in real-world situations.

Our experiments show that policies that are trained with our regularizer exhibit a simpler behavior that is more robust to mismatches in the system dynamics. We argue that these properties are highly desirable when attempting to transfer these policies to a real robot. However, we unfortunately were not able to perform experiments on a real robot.

Author is encouraged to apply their approach to a wider variety of environments.

We now additionally evaluated our approach on 5 more tasks from Metaworld. Also on these non-periodic manipulation tasks, MTC outperforms SAC and RPC in terms of average success rate, demonstrating the benefits of maximizing the total correlation on solving manipulation tasks. All in all, our experiments are quite diverse, span periodic and nonperiodic, vision-based and non-vision-based tasks, closely follow the testbeds of the most related works, and show strong performance in all of these settings.

However, optimal policies in DMC environments exhibit periodic behavior regardless of additional priors. Unclear how total correlation affects this.

While DMC tasks are characterized by periodic motions, our results in Fig.1, Fig.3, Fig.9 and Fig.10 clearly show that maximizing the total correlation improves the consistency and periodicity of behaviors compared to leading baselines. Furthermore, we are primarily interested in learning simple behavior, which does not necessarily imply periodicity, for example, for manipulation tasks.

[Comparison to literature such as TD-MPC and BYOL]

We now provide the following discussion of these two approaches in our Related Work Section: Our approach is also related to previous approaches that extract temporally consistent representations from observations by learning latent dynamics models, such as TD-MPC and BYOL. Different from these approaches only considering temporal consistency in representations of states, our total correlation objective aims to maximize the consistency among the whole trajectories of state representations and actions. As shown in our experiments (Fig.5), additionally enforcing the consistency within action sequences by learning the action prediction model improves the performance in the presence of environmental perturbations.