World Models as Reference Trajectories for Rapid Motor Adaptation

We thank you for your very positive and encouraging review. We are delighted that you recognized the novelty, simplicity, and theoretical grounding of our work.

Q1. More evaluation (Cheetah or Humanoid): Thank you for this excellent suggestion. We have now evaluated our method on Cheetah, Hopper, and Ball-in-Cup, and have new results for Humanoid under perturbation, which confirm the generality of the results obtained for the Walker environment. We will integrate these new results into the final manuscript to better showcase the method's scalability and strengthen our claims.

Q2. Clarification on Figure 5: This is a great observation that points to a key distinction. The experiment in Figure 5 is different from Figure 2: the baselines in Figure 5 were pre-trained with domain randomization on the same perturbation distribution used for testing, whereas the baselines in Figure 2 were not.

The domain-randomized TD-MPC2 agent outperforms RWM in the "OFF" phase (no perturbation) for two reasons:

1. This "average" condition matches its training distribution, making it an expert for that specific setting.
2. The agent had not fully converged during pre-training and continues to slowly fine-tune its policy for this average case.

However, the crucial test is performance during active perturbations ("ON" phase). Here, the pre-trained agent still suffers a significant performance drop, whereas RWM's online mechanism maintains higher reward. This demonstrates that pre-training on a known noise distribution is not a substitute for a true online adaptation mechanism like RWM. We will clarify this narrative in the main text.

Q3. Mismatch of the line colors in Figure 2, 5: Thank you for spotting this. We apologize for the error and will correct the figure captions to match the plots in the final version.

We thank you once again for your strong support and helpful feedback.

Thank you for response. I have carefully read the comments from other reviewers and each rebuttals.

R1. Thank you for the rebuttal

• I appreciate the authors' clarifications regarding my original concerns, all of which have been adequately addressed.
• In particular, the results in the "ON" phase demonstrate that RWM performs better than standard RL without relying on additional Sim2Real techniques such as domain randomization.
• I believe this highlights the potential of the method in settings where Sim2Real techniques are difficult to apply (e.g., real-world RL).

R2. New concern raised from other reviewers

• However, after reading the discussion and revisiting the paper, I realized that the TD-MPC2 baseline was run without planning.
• I also agree with Reviewer Mvgi that this setup may result in an unfair comparison (does not convince that RWM outperforms standard planning).
• I had initially overlooked this point and appreciate the clarification.

Given the potential impact of this issue on the empirical claims (R2), I will maintain my current score until this concern is addressed.

We thank you for your exceptionally thorough feedback. You have provided a clear roadmap for significantly improving the paper's clarity and rigor. We acknowledge that our initial presentation had shortcomings, and we will address every point.

1. On Lack of Formalism and Confusing Notation:

We sincerely apologize for the lack of clarity. We will add a new subsection to formally define the problem setting and a glossary to the appendix. To address the notational ambiguity you correctly identified, we will refer to the encoder as $\text{enc}(\cdot)$ throughout the text, distinguishing it from the error term $e_t$.

Here are new formalism sections to be added before the model design, in addition to a Glossary to be added to the appendix:

We consider a continuous control problem formulated as a Markov Decision Process (MDP), defined by the tuple $(\mathcal{S}, \mathcal{A}, P, R, \gamma)$, where $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space, $P(s'|s,a)$ is the transition probability, $R(s,a)$ is the reward function, and $\gamma$ is the discount factor. Our goal is to learn a policy $\pi(a|s)$ that maximizes the expected cumulative discounted reward. We operate in a setting where the true system dynamics, represented by the transition function $P$, can change unexpectedly from the training conditions.

Our approach leverages a learned world model operating in a latent space $\mathcal{Z}$. The model consists of: an encoder $z_t = \text{enc}(s_t)$, a latent dynamics model $\hat{z}_{t+1} = F(z_t, a_t)$, and a base policy $a_0 = \pi_0(z_t)$. The objective of our Reflexive World Model (RWM) is to learn a second, adaptive control policy, $\pi_c$, which outputs a corrective action $a_c$. The total action is $a_t = a_0 + a_c$. The objective for $\pi_c$ is to minimize the discrepancy between the world model's one-step prediction (under the base policy) and the observed outcome, thereby making the system behave as predicted by the reference model $F$ under the base policy $\pi_0$.

2. On Confusing Narrative and Experimental Design:

You raised several important points about the narrative and the experimental choices.

On Novelty and the Goal of the Method: We will clarify that the main goal is not merely to "eliminate noise." The goal is to maintain the intended behavior of the base policy ($\pi_0$) in the face of any dynamic shift that causes a deviation from the world model's prediction. Action perturbation is just one clean and measurable way to test this core principle. The novelty lies in re-purposing the world model's one-step prediction as an implicit reference trajectory for a controller, a general mechanism for adaptation that goes beyond simple noise filtering.

On Experimental Persuasiveness: You noted that our chosen perturbation, $a_{eff}=a \cdot (1+p)$, simplifies the learning objective for our controller. We agree that our method is well-suited to this type of perturbation. However, we see this as a strength of the design, not a weakness of the evaluation. This form of perturbation is a realistic model for actuator degradation or miscalibration. The fact that our method provides a principled and direct solution to this common problem is a key contribution. Furthermore, the difficulty that the strong TD-MPC2 baseline has in compensating for this "simple" perturbation demonstrates that the task is non-trivial and requires a specialized mechanism like RWM for effective, rapid adaptation.

3. On Technical Questions:

• Q1 (Forward Model $F$): This is an excellent question. The forward model $F(z,a)$ is trained as a general transition model during the pre-training phase. This means it is exposed to a wide distribution of states and actions, including those from early, suboptimal policies, not just trajectories from the final converged policy $\pi_0$. This broad training allows it to make reasonable predictions even for states that are off the optimal trajectory, which is crucial for providing a corrective signal when perturbations occur. We will clarify this in the revision.

• Q2 (Hessian): Thank you. You are correct. The Hessian of the value function, $H = \nabla^2 V$, should be negative definite near a maximum. We will correct this typo.

• Q3 (Formatting): We will fix all hyperlink and referencing issues.

• Q4 (Theorem 4.2): This is a key point. You are correct that the model $F$'s predictions will be wrong post-perturbation. However, our method does not require the absolute prediction to be correct. It relies on a weaker, more plausible assumption: that the local gradient of the model remains informative. While the perturbation shifts the dynamics, making $F(z,a)$ inaccurate, the relationship between an action change $\Delta a$ and the resulting state change $\Delta z$ is assumed to be roughly preserved. The model still "knows" that increasing a certain action will move the latent state in a particular direction. This is what the gradient $\partial F/\partial a$ captures, and as long as this gradient points in a useful corrective direction, our controller can successfully reduce the error. We will expand this explanation in the paper.

• Q5 (Algorithm 1): We will clarify that Algorithm 1 describes the online update rule applied at each timestep of an episode until its end.

We are grateful for your diligence and are confident these revisions will resolve the noted issues.

We sincerely thank you for your detailed and insightful review. We are encouraged that you found our core idea "novel and conceptually clean," with a "solid theoretical foundation." We appreciate the constructive feedback and would like to address the identified weaknesses and questions.

1. On Limited Environments and Generality Claims: You raise a valid point regarding the specificity of our claims on nonstationarity. In the revised manuscript, we will be more precise, framing our results as demonstrating robustness to dynamics arising from actuator drift and gradual miscalibration.

To better demonstrate the method's generalizability, we have also extended our experimental evaluation. We now have results on Cheetah, Hopper, Ball-in-Cup, and Humanoid, with a similar profile to the Walker simulations: stronger adaptation than the comparisons, though with a smaller effect for the Humanoid environment. Furthermore, to show the method's potential beyond actuator effects, our preliminary experiments on body mass perturbations in the Cheetah environment show the characteristic decrease in control error during perturbations. All new findings will be integrated into the final version of the paper.

2. On Baselines and Comparison to Alternative Methods: Thank you for these important points. Our goal was to test our specialized adaptation mechanism against relevant and strong baselines.

• TD-MPC2 Baseline Configuration: We thank the reviewer for this question, which allows us to clarify a critical aspect of our experimental design. Our primary goal is to compare different reactive adaptation mechanisms, as online planning is often computationally prohibitive. A key consideration is that when a policy is trained with an online planner, it learns to provide actions that serve as a good starting point for the planner, not to solve the task on its own; the policy and planner become entangled. Using such a policy without its planner during evaluation would be an unfair comparison against a crippled agent. Conversely, keeping the planner active would introduce a confounding variable, making it difficult to isolate the performance of the underlying adaptation mechanism. Therefore, to ensure a clean comparison, our "TD-MPC2" baseline uses the reactive policy head (trained without the planner) and continues to fine-tune online. Both this baseline and RWM rely on the world model for learning — TD-MPC2 for its policy and value functions, and RWM for its corrective controller. The exception is the Humanoid environment, where a reactive policy alone is insufficient to produce stable locomotion. For this challenging task only, the TD-MPC2 baseline uses its full online planner to achieve a competent level of performance. We will clarify this configuration in the revised manuscript. To provide even stronger comparisons, we will also finalize our ongoing simulations using TD-MPC2 with planning, as appropriately suggested in your feedback, and include them in the revision.

• Model Reference Adaptive Control (MRAC): Directly applying classical controllers like MRAC to high-dimensional tasks with learned, black-box dynamics models is generally not feasible. These methods require pre-defined, analytical reference models and full state observability. Furthermore, they are designed for tracking problems, typically with quadratic cost functions defined around a pre-specified reference trajectory. This is fundamentally different from the deep RL paradigm, which optimizes behavior based on general (and often complex or sparse) reward functions where the optimal trajectory itself must be discovered through interaction. A core contribution of our work is precisely bridging this gap: we introduce a mechanism that achieves MRAC-like reference tracking, but does so within a modern deep RL framework by using the world model's own predictions as an implicit, emergent reference trajectory.

• Residual Policy Learning: This brings up a useful comparison. Residual Policy Learning often starts with a classical controller as a stable base policy and adds a flexible, learned RL policy as the residual to handle complex, unmodeled dynamics. Our work can be seen as a conceptual inversion of this design. We start with a flexible, general-purpose RL policy ($\pi_0$) and add a corrective controller ($\pi_c$) that is conceptually closer to classical control, as it performs deterministic error correction based on the world model's reference signal. This architectural similarity and distinct design choice will be clarified in our related work section.

• Meta-RL: Our work addresses continuous online adaptation to unforeseen dynamic shifts, whereas Meta-RL learns an adaptation strategy from a pre-defined distribution of training environments. We therefore focused our comparison on online fine-tuning, which more directly aligns with our core research question.

3. On Experimental Clarity and Organization: We agree that the experimental section can be improved. In our revision, we will restructure Section 5 to first clearly define the experimental questions, then introduce all environments, and finally present the results systematically. We will also clarify Table 1 and properly introduce the point-mass system.

We thank you again for your time and effort. We are confident that these clarifications will significantly strengthen the manuscript.