ResWM: Residual-Action World Model for Visual RL

1. Problem Statement

Reinforcement Learning (RL) agents operating in high-dimensional visual environments (e.g., robotics) face significant challenges in learning predictive world models. Traditional Model-Based RL (MBRL) frameworks, such as Dreamer and TD-MPC, typically condition future state predictions on absolute actions ($a_t$). This approach introduces several critical issues:

• Optimization Instability: The distribution of optimal absolute actions is task-dependent and non-stationary, making it difficult for the model to learn a consistent policy.
• High Variance & Oscillation: Directly predicting absolute actions often leads to erratic, high-frequency control signals ("chattering"), which are physically unrealistic and unsafe for real-world robotic systems.
• Inefficient Search Space: The agent must search the entire global action space at every step, ignoring the inherent temporal smoothness of physical dynamics.
• Redundant Perception: Standard encoders often process static frames independently, failing to explicitly capture the temporal deltas that drive physical changes.

2. Methodology: Residual-Action World Model (ResWM)

The authors propose ResWM, a framework that reformulates the control variable from absolute actions to residual actions (incremental adjustments) and couples this with a dynamics-aware observation encoder.

A. Residual Action Policy (Smoothness Prior)

Instead of predicting the absolute action $a_t$, the policy predicts an incremental adjustment $\delta a_t$ relative to the previous action $a_{t-1}$.

• Formulation: $a_t = \tanh(a_{t-1} + \delta a_t)$, where $\delta a_t \sim \pi_\theta(\cdot | z_t, a_{t-1})$.
• Benefit: This embeds a strong temporal smoothness prior into the action space. It transforms the optimization problem from a chaotic global search to a localized refinement around the previous state, drastically reducing the effective search space and stabilizing long-horizon planning.

B. Observation Difference Encoder (ODL)

To align the perception module with the residual action paradigm, the authors introduce the Observation Difference Encoder.

• Mechanism: Instead of encoding static frames, the ODL explicitly models the difference between consecutive observations ($o_t$ and $o_{t-1}$).
• Architecture: $z_t = \text{LN}(\text{FC}(f(o_t) - f'(o_{t-1})))$, where $f$ and $f'$ are CNN encoders.
• Benefit: This acts as a temporal filter, suppressing static background noise and extracting only the dynamic, task-relevant changes (e.g., moving joints). This creates a latent representation $z_t$ that is causally coupled with the residual action $\delta a_t$.

C. Latent Dynamics & Training

ResWM integrates these components into a Recurrent State-Space Model (RSSM) (similar to Dreamer):

• Transition Model: $s_{t+1} \sim g_\phi(s_t, \delta a_t)$. The dynamics are conditioned on the residual action, ensuring consistency between the imagined rollouts and the policy's behavior.
• Objective: The model is trained end-to-end using a joint loss function comprising:
  1. Model Loss: VAE ELBO (image reconstruction, reward prediction, KL divergence).
  2. Policy Loss: Actor-Critic optimization on imagined rollouts within the residual space.
  3. Regularization: A KL penalty encourages $\delta a_t$ to stay near a zero-mean Gaussian (preventing erratic jumps), and an optional energy penalty minimizes control effort ($\|\delta a_t\|^2$).

3. Key Contributions

1. Residual-Action World Model (ResWM): A novel MBRL framework that reparameterizes the action space to enforce temporal smoothness, significantly alleviating learning complexity and improving stability.
2. Observation Difference Encoder (ODL): An architecture that explicitly encodes temporal deltas between frames, creating a compact, dynamics-aware latent space that naturally aligns with residual action prediction.
3. Plug-and-Play Integration: The framework requires minimal modifications to existing Dreamer-style architectures and introduces no new hyperparameters, making it highly adaptable.
4. Empirical Validation: Comprehensive evidence showing that ResWM outperforms state-of-the-art baselines in sample efficiency, asymptotic performance, and control smoothness.

4. Experimental Results

The authors evaluated ResWM on the DeepMind Control Suite (DMControl) and Atari benchmarks.

• DMControl (Standard Tasks):

  • At 100K steps, ResWM achieved an average score of 828.7, outperforming the strong baseline DeepRAD (695.1) and ResAct (822.1).
  • At 500K steps, the average score rose to 925.0, surpassing all competitors (e.g., DeepRAD at 890.8).
  • ResWM achieved near-saturation performance on tasks like Reacher, Easy (986) and Ball in cup, Catch (978).

• DMControl (Hard Tasks):

  • On complex tasks (e.g., Quadruped, Walk), ResWM maintained a lead at 1M steps with an average score of 644.8, significantly beating ResAct (630.2) and TACO (333.8).

• Atari Benchmark:

  • ResWM achieved a normalized mean score of 0.96, outperforming ResAct (0.86) and TACO (0.76), demonstrating strong generalization to discrete, high-dimensional pixel environments.

• Ablation Studies:

  • Removing the Residual Policy (V1) caused catastrophic failure on hard tasks (e.g., Hopper, Hop), proving it is the most critical component.
  • Removing ODL (V2) resulted in the second-largest performance drop, highlighting the importance of dynamics-aware perception.
  • Qualitative Analysis: Attention maps showed ResWM focuses sharply on critical joints and effectors, whereas baselines (DeepRAD) exhibited diffuse attention across entire limb contours.

5. Significance

• Bridging Theory and Practice: ResWM addresses the disconnect between algorithmic RL and the physical requirements of robotics (smoothness, energy efficiency, safety).
• Stability & Safety: By reducing action variance and chattering, ResWM produces trajectories suitable for real-world deployment where mechanical wear and energy consumption are critical constraints.
• Efficiency: The framework achieves superior sample efficiency, meaning agents can learn complex tasks with fewer interactions, a crucial factor for real-world robotics where data collection is costly.
• General Principle: The paper suggests that modeling residual changes rather than absolute states is a fundamental principle for improving continuous control in visual RL, applicable beyond the specific architectures tested.

Limitations: The authors note that the residual prior might slightly constrain the agent's ability to react instantly to abrupt, high-frequency shocks (e.g., sudden collisions) and introduces minor computational overhead during the encoding phase due to processing consecutive frames.