Maximum Total Correlation Reinforcement Learning

• The concept of total correlation is quite similar to mutual information, which has been extensively explored in both single-agent and multi-agent RL. Therefore, the novelty is limited, with the primary distinction being that prior work considers state-action pairs, while total correlation extends this to entire trajectories.
• Although the method was evaluated on 8 tasks in the DeepMind Control (DMC) suite, these tasks are relatively simple and inherently periodic. For such tasks, structurally designed policies or dynamic movement primitives (DMP) could potentially be more efficient. Therefore, it is essential to assess the generalizability of maximizing total correlation in more complex tasks, such as robotic grasping or kitchen environments.

The main claims are not sufficiently supported.

1. While this paper emphasizes the unified framework of state-consistency and action-consistency, the authors did not investigate whether the unified framework is required for their results. I strongly recommend conducting an ablation study using either state-consistency or action-consistency to support claims 1), 2), and 3).
2. Regarding claims 1) and 2), although this paper includes the state-of-the-art algorithm LZ-SAC, that paper also includes another algorithm called SPAC that imposes simplicity on action sequences. LZ-SAC paper reports that while LZ-SAC achieves better final performance than SPAC, SPAC shows better results in robustness and information efficiency. To fully support the claims 1) and 2), I recommend including SPAC as an additional baseline.
3. Regarding claim 1), while both RPC and LZ-SAC are expected to show robustness to dynamics changes, Figure 2 shows that they perform worse than SAC, which gives inconsistency of experiments.
4. Some works such as SPAC use transformers for their action sequence compressor, while other works such as this work use LSTM. Could the authors show that the merit of their work comes from the proposed unified framework, not the choice of the type of compressors?

• My primary concern is the definition of the regularized reward function $r^*(s_t,a_t,s_{t+1})$. This reward depends on the entire history of the episode through the dynamics model $q_\mu(z_{t+1}|z_{1:t},a_{1:t})$. It should thus correctly be written $r^*(s_{a:t},a_{1:t},s_{t+1})$ and is no longer Markovian. However, this non-markovianness or its implications are not discussed anywhere. I would expect this to cause severe instabilities in training, but maybe that's not the case? At the very least this should be discussed and I would like to see some visualization of the training dynamics in the form of a training step vs reward plot.

• The experimental validation is somewhat limited. Hyperparameter optimization is not performed in a structured way, both for the baselines and the proposed method.

• The introduction motivates the method through combatting spurious correlation but there are not distractors etc. present in the experiments that could lead to spurious correlation beyond a noisy state observation. Likewise, the utilization of a learned state representation $z$ is motivated by the noisy TV problem, but this also isn't directly investigated in the experiments.

• An ablation of only using the latent-dynamics model but not the action model would also be useful to see the contribution of the different parts of the method. This would be somewhere between the proposed method and RPC and would thus be very interesting to see.

• This method adds additional dynamic models that are trained along with the the policy, so naturally the training time should be longer.

Structured into claims:

• Claim 3: The method reduces spurious correlation and focuses only task-relevant state information by using a learned state encoding. Evidence: Only indirect evidence is given through experiments with observation noise. I don't think this is sufficient to support the claim and experiments with explicit distractors would be more convincing.
• Claim 4: Maximum Total Corrlation SAC as a well motivated RL method. Evidence: The derivation hides the non-markovianness/nonstationarity of the reward function by incorrectly listing its arguments. This is my main concern with this paper and not discussed at all. The experiments indicate that it's practically not a big issue but I would like this to be elucidated more.

Minor comments:

• Line 278: missing closing bracket for expectation
• Line 671: $q(a_{t+1}...$ should be $q_\mathcal{X}(a_{t+1}...$
• It would also be nice to discuss the connection to works that directly learn periodic controllers such as [1], and maybe even the connection to periodic locomotion controllers which are popular in classical robotics.

[1] Raffin, Antonin, et al. "An Open-Loop Baseline for Reinforcement Learning Locomotion Tasks." RLC 2024

1. The fundamental motivation of the paper, "simplicity is a powerful inductive bias", is not well supported. The authors support it by enumerating examples in related problems, e.g. "regularization is used for simpler policies, data augmentation for simpler representations, and sparse reward functions for simpler objectives." But this is quite vague. It might very well be that in RL, overoptimizing for simplicity leads to suboptimality in complex tasks.
2. In Section 4.1, the authors motivate the use of latent states by stating that it is more robust to distractions from task-irrelevant factors (e.g. noisy TV). This claim again lacks support.
3. The key differences between MTC and RPC [1] are (1) MTC considers states and actions while RPC only considers states, and (2) MTC conditions the dynamics and action prediction models on the history, whereas RPC only conditions on the current state. I'm not sure why (1) would conceptually improve upon RPC (see question below). (2) introduces a confounding variable that is orthogonal to the claim, since we can replace the model in RPC with a recurrent model and hope to see better performance.
4. While the goal is to induce simplicity in the policy's behavior, the method requires learning three additional models to estimate an auxiliary objective, which is quite cumbersome.
5. The rest of experiments are unconvincing. While the main results (Table 1) show MTC performing on par or slightly better than baselines, the improvement is not significant. Moreover, these results are in a narrow domain that prefers repetitive motion (e.g. Hopper, Cheetah, Walker are all locomotion tasks). I'm not sure if the same holds for more complex tasks. The robustness results in Figure 2 don't imply MTC being more robust than baselines, because SAC shows the same growth trend, just starting from a lower initial value. What I hope to see is that MTC performance degrades at a slower rate than SAC, instead of having a constant offset.

Overall, I think this paper can be significantly strengthened by (1) grounding it in a concrete problem setting like sim2real transfer instead of a lofty desideratum (i.e. simplicity), and (2) either removing the unwarranted claims or adding experiments to support them. I strongly recommend the authors to revise and resubmit.

List of Claims (cont'd)

Claim 2: simplicity is desirable for RL.

• Assumption: none
• Evidence: slightly better empirical performance of MTC over baselines on DM control tasks.
• Statement: I don't think there is enough theoretical or empirical evidence to support this claim. Given that the DM control tasks prefer periodic behavior, it might be that simplicity is particularly useful in this domain.

Claim 3: using latent states results in less distraction from task-irrelevant factors.

• Assumption: the environment consists of task-irrelevant distractors, modeling which hinder the policy's performance.
• Evidence: none
• Statement: to support this claim, the authors need to compare their method to a baseline that uses raw observations on environments with distractors.

Claim 4: MTC improves upon RPC by considering states and actions, and conditioning the models on the history.

• Assumption: none
• Evidence: slightly better empirical performance than RPC.
• Statement: I'm not sure "why" these changes would improve upon RPC. The sequence of states should determine the sequence of actions to a large extent. Conditioning on the history is a choice of variational approximation that can also be applied to RPC, and the paper doesn't compare to RPC + recurrent models.

[1] Benjamin Eysenbach, Ruslan Salakhutdinov, Sergey Levine. Robust Predictable Control. NeurIPS 2021.