RobustZero: Enhancing MuZero Reinforcement Learning Robustness to State Perturbations

We appreciate your positive and valuable comments.

$**Response to Methods and Evaluation Criteria:**$ 1) Regarding the projector and predictor networks, please refer to the response to Q1; 2) Regarding the bold entries, we have revised all tables (see https://anonymous.4open.science/r/RobustZero-SupportMaterials-8512/Tables/Tables%201-5.png); 3) We have added the learning curves (see https://anonymous.4open.science/r/RobustZero-SupportMaterials-8512/Figures/Figure%20S-1.png).

$**Response to Experimental Designs or Analyses:**$ We have revised all tables according to your suggestion.

$**Response to Other Comments or Suggestions:**$ We would like to address your comments as follows:

1. MuZero-class methods learn abstract environment models, i.e., learned models, and employ MCTS for planning. At each time step, a real state is mapped to an abstract initial state by using the representation network. Then, the dynamics network and prediction network work in conjunction with MCTS to predict future evolution of these abstract states. This framework inherently involves a prediction process. The compound errors mean the accumulation of inaccuracies in these predicted abstract states. To clarify this point, we will revise the statement to ‘Even a small “reality gap” can compound errors in the predicted abstract states within MuZero-class methods, causing reduced rewards and potentially harmful decision-making’.

2. We will supplement the following missing information: i) $\rho_0$ is the initial state distribution; ii) $\mathcal{P}$ is a deterministic transition function; and 3) the state space is continuous.

3. We will use “state” uniformly.

4. We will add the following statement. The sequence of real actions from the sampled trajectory is used to compute $r_t^k$.

5. We will revise the statement to “the goal is to maximize the expected return.”

6. We have carefully double-checked the MuZero algorithm, and make sure that the reward $r_t^k$ is the output of the dynamics network $\mathcal{G}$.

7. We will replace “environment model” with “MDP”.

8. We will revise the statement to “$\widetilde{t}$ and $\hat{t}$ denote data at time t from a trajectory generated under worst-case and random perturbations, respectively.”

9. We have revised Fig. 1 by adding the mathematical denotations of the networks (see https://anonymous.4open.science/r/RobustZero-SupportMaterials-8512/Figures/Figure%201.png).

10. We have defined $\mathcal{D}\_{*WC*}$ and $\mathcal{D}\_{*RC*}$ in Fig. 1 (see the link above).

11. We will revise the expressions: i) From $\mathcal{D}\_{*WC*}\sim\mathcal{B}$ to $\mathcal{B}\sim\mathcal{D}\_{*WC*}$; ii) From $\mathcal{D}\_{*RC*}\sim\mathcal{B}$ to $\mathcal{B}\sim\mathcal{D}\_{*RC*}$; and iii) From “B is the batch size” to “B is the batch”.

12. We will add the missing symbol for Eq. 10.

$**Response to Q1:**$ Model collapse refers to a degenerate solution in which the encoder produces identical or nearly identical outputs for all inputs, resulting in trivial and non-informative representations. If we omit the predictor and directly enforce similarity between the outputs of two branches (e.g., using mean squared error or cosine similarity), the model may exploit a shortcut: “I might as well just output a constant vector, since all inputs yield the same output, the loss will be minimized.” This leads to model collapse. When collapse occurs, all states—regardless of whether they are perturbed or not—are mapped to the same or very similar initial hidden state, effectively losing the ability to distinguish between different inputs. To avoid model collapse, one effective strategy is to introduce an asymmetric network architecture, where a predictor is added to one branch, while the other branch remains without it and has its gradient flow stopped. This asymmetry prevents both branches from trivially converging to the same constant output. The branch with the predictor learns to align its output with the target branch, which acts as a stable reference point. Because the target branch does not receive gradients, it cannot adjust itself to match the predictor’s potentially trivial solution, thus breaking the symmetry that often leads to collapse. This design encourages the network to learn meaningful representations, resulting in stabilization of optimization and a reduction in collapsing solutions.

$**Response to Q2:**$ The reasons are as follows. First, as shown in Eqs. 3-4, both the worst-case and random-case loss terms already incorporate information from perturbed and unperturbed states due to the use of a contrastive loss function. This contrastive loss encourages the learned policy to remain consistent before and after state perturbations. Therefore, an additional MuZero loss term specifically targeting unperturbed states is redundant. Second, an additional loss term will increase the computational cost.

We appreciate your valuable comments and recognition of our contributions.

$**Response to Experimental Designs or Analyses:**$ Please refer to the responses to Q1-Q3.

$**Response to Essential References Not Discussed:**$ We will add the two references as follows. The recent studies [1-2] introduce the diffusion models to enhance robustness to state perturbations. Specifically, a diffusion model-based predictor is proposed [1] for offline RL to recover the actual states against state perturbations. A belief-enriched pessimistic Q-learning method is proposed [2] by using diffusion model to purify observed states.

$**Response to Q1:**$ Following your suggestion, we have studied RobustZero and all baselines on five Mujoco environments, including Hopper, Walker2d, HalfCheetah, Ant and Humanoid. The results are provided in Table S-1 (see https://anonymous.4open.science/r/RobustZero-SupportMaterials-8512/Tables/Table%20S-1.png). The results show consistently that RobustZero outperforms all baselines at defending against state perturbations. This provides evidence that RobustZero performs well in general and challenging Mujoco environments.

$**Response to Q2:**$ We fully agree with the reviewer that PA-AD (referred to as PA-ATLA-PPO in our paper) is currently considered the strongest attack method. However, as mentioned in Section 2.2, it is a white-box attack, requiring access to the internal parameters of the victim model. In contrast, our work focuses on black-box attack settings, where such access is not available. Therefore, we do not include PA-ATLA-PPO as a baseline in our experiments.

$**Response to Q3:**$ Since SA-PPO is a white-box attack strategy, we do not use it as a baseline.

$**Response to Q4 and W1:**$ Our method involves four parameters: $w1$, $w2$, $\lambda_1$, and $\lambda_2$. We clarify their design and selection: 1) One of our contributions is the development of an adaptive mechanism to adjust $w1$ and $w2$, eliminating the need for ad-hoc adjustments. Note that in Appendix C.3, we show the selection of $w2$ across different environments. This is used in ablation studies, where we intentionally fix $w2$ and find best value to isolate its effect and demonstrate the advantage of the adaptive mechanism; 2) $\lambda_1$ adjusts the trade-off between robustness to worst-case perturbations and random perturbations. A larger $\lambda_1$ emphasizes worst-case robustness, while a smaller value favors random-case robustness. When both are considered equally important, setting $\lambda_1=1$ (as done in our experiments) is a reasonable and effective default. Thus, $\lambda_1$ does not require tuning unless specific emphasis is desired; and 3) Among the four parameters, $\lambda_2$ is the only one that requires manual tuning. To keep the selection process simple, we employ a standard grid search. Appendix C.7 analyzes the effect of different $\lambda_2$ values, which is not the selection process. In summary, we only need to search for the best setting of $\lambda_2$ for a new environment. Therefore, the hyperparameter selection is not complex.

$**Response to Q5:**$ RobustZero does not know the attack budget $\epsilon$ during training. We will clarify this statement in the final paper.

$**Response to Q6:**$ Literature [3] is the first study to design a defense strategy against state perturbations specifically for offline RL. The strength of offline RL lies in its ability to train agents solely from a fixed dataset, without requiring any interaction with the environment. However, this limits its exploration capability, making it more challenging to learn optimal policies. That said, [3] demonstrates strong robustness to state perturbations when compared to other offline RL methods that lack defense strategies. In contrast, RobustZero and all baseline methods evaluated in our work are online RL methods, where the agent learns by actively interacting with the environment. Generally, online RL methods are capable of achieving higher rewards than offline RL methods due to their adaptive exploration. Given the fundamental differences in training protocols and assumptions between offline and online RL, we do not include [3] as a baseline, as it would not constitute a fair comparison under the same experimental setup.

$**Response to Q7:**$ In Appendix C.6, we provided comparison analysis for the training time and sampling time. Since the sampling time closely approximates the testing time, we initially omitted testing time results in the submitted paper. Now, we have added the testing time per step (TeT) and updated Table 6 accordingly (refer to https://anonymous.4open.science/r/RobustZero-SupportMaterials-8512/Tables/Table%206.png). As shown in the update Table 6, the testing time is similar to the sampling time, the previous analysis of sampling time remains applicable. For a comprehensive comparison, please refer to Appendix C.6.

We appreciate your positive and valuable comments.

$**Response to Claims and Evidence:**$ Following your comments, we have analyzed the relationship between the number of environment samples and the natural, worst-case, and random-case rewards for RobustZero and the two robust model-free baselines: ATLA-PPO and PROTECTED. The results are presented in Table S-2 (see https://anonymous.4open.science/r/RobustZero-SupportMaterials-8512/Tables/Table%20S-2.png). Please note that the number of samples per episode in RobustZero differs from those used in the two baselines (refer to Table 6 in Appendix C.6). Therefore, while the numbers of samples reported in Table S-2 are similar across methods, they are not exactly the same. From Table S-2, by using similar samples, RobustZero achieves higher rewards compared to ATLA-PPO and PROTECTED, demonstrating its superior sample efficiency.

$**Response to Method and Evaluation Criteria:**$ Following your suggestion, we have studied RobustZero and all baselines on five Mujoco environments, including Hopper, Walker2d, HalfCheetah, Ant and Humanoid. The results are provided in Table S-1 (see https://anonymous.4open.science/r/RobustZero-SupportMaterials-8512/Tables/Table%20S-1.png). The results show consistently that RobustZero outperforms all baselines at defending against state perturbations. This provides evidence that RobustZero also performs well in general and challenging Mujoco environments.

$**Response to Experimental Designs or Analyses:**$ To ensure a fair comparison, we follow prior works, i.e., ATLA-PPO and PROTECTED by using ATLA-PPO to obtain a worst-case perturbation policy that is applied to all methods. Regarding robust RL methods, e.g., RobustZero, ATLA-PPO and PROTECTED, the agent is trained by using non-perturbed states and perturbed states information. This enables obtaining consistent policies before and after state perturbations. Regarding S-MuZero-worst, it is an extended baseline, representing S-MuZero trained under the worst-case perturbation policy. Importantly, S-MuZero-worst does not incorporate any robustness mechanism. It is trained solely on perturbed states, without any information about non-perturbed states. This causes S-MuZero-worst to be over-trained under the worst-case perturbation policy, making it highly specialized and adapted to that particular perturbation pattern. As a result, its performance under the worst-case perturbation policy unusually appears higher than its performance under random or no perturbations as observed in Table 1. This explains why its worst-case reward may appear higher than its random-case reward. Note that S-MuZero-worst is deliberately designed to showcase the effect of over-training under worst-case perturbations, regardless of its natural and other performance. Except for S-MuZero-worst, the worst-case rewards are lower than the corresponding natural and random-case rewards.

$**Response to Essential References Not Discussed**:$ We would like to address your comments as follows: 1) Following your suggestion, we will add the two references. Specifically, a robust adversarial reinforcement learning method is proposed [1] to jointly train an agent and an adversary, where the agent aims to accomplish the primary task objectives while learning to remain robust against disturbances introduced by the adversary. A recent study [2] proposes the use of variational optimization over worst-case adversary distributions, rather than a single adversary, and trains an agent to maximize the lower quantile of returns to mitigate over-optimism; and 2) In our paper, a worst-case state perturbation refers to an adversarial modification of the agent’s observed state that is carefully crafted to minimize its expected return. While, a random-case state perturbation refers to a stochastic disturbance applied to the agent’s observed state. Formal definitions of both perturbation policies are provided in Definitions 4.2 and 4.3 of our paper (see pages 3-4). Additionally, we would like to note that for reference [2], only the abstract is currently available, and the full version has not yet been released online. As a result, we are unable to determine the specific settings used in that work regarding worst-case and random-case perturbations.

$**Response to Other Comments or Suggestions:**$ We will change “row” to “column”.

$**Response to Questions for Authors:**$ S-MuZero can solve the tasks with continuous action spaces. We will add this statement in Section 3.3.

We appreciate your comments and our responses are detailed below.

$**Response to Method and Evaluation Criteria and Q1**$: We would like to address your comments as follows:

1. We have studied RobustZero and all baselines on five Mujoco environments, including Hopper, Walker2d, HalfCheetah, Ant and Humanoid. The results are provided in Table S-1 (see https://anonymous.4open.science/r/RobustZero-SupportMaterials-8512/Tables/Table%20S-1.png). RobustZero consistently outperforms all baselines at defending against state perturbations. This provides evidence that RobustZero also performs well in general and challenging Mujoco environments.
2. Following ATLA-PPO and PROTECTED, we adopt an $l_{p}$-norm perturbation model to generate small, semantically invariant perturbations. However, in board games, the states are discrete and highly structured, and even small $l_{p}$-norm perturbations can result in invalid or illegal states that violate game rules. Therefore, this perturbation model is not suitable for board games, and we do not evaluate all methods in such environments.
3. We will add the following explanations for the energy and transportation tasks. The three transportation environments support the testing of autonomous driving tasks. Therein, an autonomous driving car interacts with other vehicles to navigate different scenarios: i) Highway$-$Drive fast, avoid collisions, and stay in the right-most lane; ii) Intersection$-$Cross safely, follow traffic rules, and keep a steady speed; and iii) Racetrack$-$Finish quickly while staying on track and driving smoothly. The action space of an autonomous driving car is two-dimensional. The three energy environments support the testing of voltage control tasks with the objective of minimizing the total cost of voltage violations, control errors, and power losses, while meeting both networked and device constraints. The action spaces of IEEE 34-bus, IEEE 123-bus and IEEE 8500-node are 10-dimensional (8 continuous and 2 discrete), 15-dimensional (11 continuous and 4 discrete), and 32-dimensional (22 continuous and 10 discrete), respectively. Thus, these energy environments are complex and high-dimensional.

$**Response to Experimental Designs or Analyses**$: We have added these training curves (see https://anonymous.4open.science/r/RobustZero-SupportMaterials-8512/Figures/Figure%20S-1.png).

$**Response to Q2**$: Our method involves four parameters: $w1$, $w2$, $\lambda_1$, and $\lambda_2$. We clarify their design and selection: 1) One of our contributions is the development of an adaptive mechanism to adjust $w1$ and $w2$, eliminating the need for ad-hoc adjustments. Note that in Appendix C.3, we show the selection of $w2$ across different environments. This is used in ablation studies, where we intentionally fix $w2$ and find best value to isolate its effect and demonstrate the advantage of the adaptive mechanism; 2) $\lambda_1$ adjusts the trade-off between robustness to worst-case perturbations and random perturbations. A larger $\lambda_1$ emphasizes worst-case robustness, while a smaller value favors random-case robustness. When both are considered equally important, setting $\lambda_1=1$ (as done in our experiments) is a reasonable and effective default. Thus, $\lambda_1$ does not require tuning unless specific emphasis is desired; and 3) Among the four parameters, $\lambda_2$ is the only one that requires manual tuning. To keep the selection process simple, we employ a standard grid search. Appendix C.7 analyzes the effect of different $\lambda_2$ values, which is not the selection process. In summary, our method does not rely on ad-hoc adjustments. The adaptive mechanism removes the need to tune $w1$ and $w2$. $\lambda_1$ can be fixed to a default value, e.g., 1. Only $\lambda_2$ requires simple tuning via grid search. Therefore, the hyperparameter selection is not complex.

$**Response to Essential References Not Discussed**$: We will add the following statements. The contrastive representation learning has been used to improve the sample efficiency [A, B] and enhance the representation ability of states [C], rewards [D], and value functions [E]. Among them, one notable work is EfficientZero [B], which significantly improves the sample efficiency of the MuZero method while maintaining superior performance. It achieves this by using contrastive representation learning to build a consistent environment model and using the learned model to correct off-policy value targets. Different from these studies, we aim to leverage contrastive representation learning to improve the robustness of MuZero-class methods to state perturbations.

[A] Data-efficient reinforcement learning with self-predictive representations

[B] Mastering Atari games with limited data

[C] Planning with goal-conditioned policies

[D] Beyond reward: Offline preference-guided policy optimization

[E] Contrastive learning as goal-conditioned reinforcement learning