PAPER DEEP DIVE
DynaWM: Dynamics-Aware Distillation with World Model and Momentum Targets for Smooth Locomotion over Continuous Stairs
Recent advances in control have enabled bipedal-wheeled robots to traverse slopes and single-step obstacles, yet long staircase traversal remains challenging as current teacher-student frameworks suffer from weakened dynamics-aware representations and incomplete terrain geometry encoding. To bridge this gap, we propose DynaWM, a dynamics-aware representation learning framework. To enhance terrain encoding capability and enable transparent assessment, we introduce a world model as a regularizer to enforce forward-dynamics awareness, preserving comprehensive terrain geometry while facilitating hierarchical encoding visualization. To stabilize knowledge transfer, we employ a momentum target encoder to provide consistent distillation targets, preventing dimensional collapse from non-stationary teacher updates. Evaluation of the learned representations through Principal Component Analysis (PCA) visualization and quantitative metrics reveals that our encoder hierarchically captures terrain geometry with higher terrain encoding capability, leading to enhanced terrain adaptability and motion smoothness. Experimental results in simulation and real hardware demonstrate that our method achieves superior terrain adaptability and motion smoothness, enabling bipedal-wheeled robots to overcome diverse continuous stairs, as shown in Fig. 1.
DynaWM: Dynamics-Aware Distillation with World Model and Momentum Targets
Paper: DynaWM: Dynamics-Aware Distillation with World Model and Momentum Targets for Smooth Locomotion over Continuous Stairs | Authors: Haidong Hou, Zhangguo Yu, Hengbo Qi, Jianlin Zhang | Institution: Beijing Institute of Technology | Link: https://arxiv.org/abs/2606.24089
One-Sentence Summary
DynaWM introduces a world model as forward-dynamics regularizer and a momentum target encoder for stable distillation, solving weakened dynamics-aware representations and dimensional collapse in teacher-student frameworks for bipedal-wheeled robot continuous stair traversal, achieving superior terrain adaptability and motion smoothness.
Background and Motivation
Bipedal-wheeled robots combine the energy efficiency of wheeled locomotion with the terrain adaptability of legged systems. While MPC and DRL advances have enabled traversal of slopes, grasslands, and single-step obstacles, long staircase traversal remains challenging due to weakened dynamics-aware representations and incomplete terrain geometry encoding in current teacher-student frameworks.
Privileged learning paradigms (teacher-student distillation) train a teacher with exteroceptive information (terrain elevation maps, contact forces) then distill knowledge to a student using only proprioceptive observations. However, three main issues exist: (1) teacher encoders optimized solely through policy gradients weaken dynamics-aware representations, encoding only terrain features directly relevant to immediate rewards; (2) concurrent teacher-student updates cause rapid teacher representation changes leading to student dimensional collapse; (3) the latent space is an opaque black box where terrain height modeling capability cannot be verified.
Figure 1: Models trained with our method evaluated across diverse outdoor stairs.
Method
DynaWM comprises three interconnected modules: World Model Learning (forward-dynamics regularization), Concurrent Teacher-Student training, and BYOL momentum target distillation.
Figure 2: DynaWM framework. World model regularizes teacher encoder; momentum target encoder provides stable distillation targets.
World Model for Dynamics-Aware Representation
The teacher encoder $E_{\theta^t}$ serves two functions: generating latent variables $z_t^t = E_{\theta^t}(s_t^t)$ for forward state prediction, and providing terrain features to the policy. The world model $W_\omega$ takes the encoder output $z_t^t$ and current privileged observation $s_t^t$ to predict the next privileged state:
$$\hat{s}_{t+1}^t = W_\omega(z_t^t, s_t^t)$$
Training the world model propagates dynamics-aware gradients to the teacher encoder by minimizing prediction residual:
$$\mathcal{L}_{pred}(\theta^t, \omega) = \mathbb{E}_{(s_t, s_{t+1}) \sim \mathcal{B}}\left[\left\|W_\omega(z_t^t, s_t^t) - s_{t+1}\right\|^2\right]$$
The gradient $\nabla_{\theta^t} \mathcal{L}_{pred}$ flows back to the teacher encoder, compelling it to extract essential dynamical variables. The world model is implemented as a deep residual network where each block follows:
$$\mathbf{h}_{i+1} = \mathbf{h}_i + \mathcal{F}_i(\mathbf{h}_i)$$
with $\mathcal{F}_i$ learned through four successive Dense-LayerNorm-Swish units. Skip connections facilitate stable gradient flow for capturing long-range temporal dependencies.
Momentum Target Encoder for Stable Distillation
Inspired by BYOL, a momentum target encoder addresses teacher non-stationarity. Two sets of student encoder parameters are maintained: the momentum target network $E_{\hat{\theta}^s}$ using EMA to gradually follow the online network $E_{\theta^s}$:
$$\hat{\theta}_{t+1}^s = \tau \hat{\theta}_t^s + (1-\tau) \theta_t^s$$
where $\tau \in [0.9, 0.99]$. A predictor head $q_\xi$ introduces asymmetry, preventing trivial collapse. The distillation objective uses normalized MSE:
$$\mathcal{L}_{moment} = \left\|\frac{q_\xi(z_t^s)}{\|q_\xi(z_t^s)\|_2} - \frac{\hat{z}_t^s}{\|\hat{z}_t^s\|_2}\right\|_2^2$$
An optional alignment loss $\mathcal{L}_{mse} = \|z_t^s - z_t^t\|_2^2$ can be added, with total student objective $\mathcal{L}_s = \mathcal{L}_{pred} + \lambda \mathcal{L}_{mse}$.
Reward Design
Beyond basic velocity tracking and torque limits, joint range constraint $r^{hr}$, wheelbase distance $r^{fd}$, base position bound $r^{pb}$, and leg balance $r^{lb}$ rewards are introduced. The hip rate reward $r^{hr} = \|\max(0, \mathbf{q}^{hip} - 0.05)\|_2$ with weight 0.1 encourages using knee and ankle joints rather than over-relying on hips.
flowchart TB
A["Privileged obs st
terrain map+contact force"] --> B["Teacher Encoder E_theta_t"]
B --> C["Latent z_t_t"]
C --> D["World Model W_omega
predict next state"]
D --> E["Prediction loss L_pred
regularize teacher encoder"]
C --> F["PPO Policy Training
Teacher Actor-Critic"]
G["Proprioceptive history ot
joints+IMU"] --> H["Student Encoder E_theta_s"]
H --> I["Online z_t_s"]
I --> J["Predictor q_xi"]
H --> K["Momentum Target Encoder
EMA update"]
K --> L["Target z_hat_t_s"]
J --> M["Momentum loss L_moment
normalized MSE"]
L --> M
E --> N["Gradient backprop
dynamics awareness"]
M --> H
Experimental Results
Training uses IsaacGym with 4000 parallel environments (3000 teacher + 1000 student), policy at 200Hz training/50Hz update, hardware at 500Hz. Training terrains include flat ground, slopes (≤30°), pyramid steps (≤0.18m), discrete uneven terrain (≤0.16m), and rough surfaces (≤0.03m).
| Reward Term | Definition | Weight |
|---|---|---|
| Linear vel. tracking | $\exp(-8.3\|\mathbf{v}_{xy} - \mathbf{v}_{xy}^{cmd}\|_2^2)$ | 7.0 |
| Angular vel. tracking | $\exp(-8.3\|\boldsymbol{\omega}_z - \boldsymbol{\omega}_z^{cmd}\|_2^2)$ | 4.0 |
| Orientation | $\|\mathbf{g}_{xy}\|_2$ | -10.0 |
| Hip rate | $\|\max(0, \mathbf{q}^{hip} - 0.05)\|_2$ | 0.1 |
| Leg position bias | $\|\mathbf{q}^{left} - \mathbf{q}^{right}\|_2$ | -0.8 |
| Joint position bias | $\|\mathbf{q} - \mathbf{q}^{def}\|_2$ | -1.2 |
Table 1: Partial reward terms and weights. Hip rate reward encourages knee/ankle usage.
PCA visualization shows that the teacher encoder with world model regularization exhibits clear terrain height stratification, while without the world model representations are entangled. The complete student framework successfully replicates the teacher's structured manifold, while without world model the student suffers dimensional collapse, and without momentum target the student fails to encode terrain properly.
| Ablation Variant | Pred. Loss | PC-Correlation | Terrain Adapt. |
|---|---|---|---|
| DynaWM (full) | Lowest | Highest | Strongest |
| w/o World Model | — | Low | Weak |
| w/o Momentum Target | — | Dim. collapse | Failed |
| MLP World Model | Higher | Medium | Medium |
| ResNet World Model | Medium | Higher | Good |
Table 2: Ablation study. Dense-LayerNorm-Swish world model achieves best prediction loss and policy performance.
Figure 3: PCA visualization. (a) Teacher with world model shows terrain height stratification; (d) Student without world model suffers dimensional collapse.
Figure 4: Training curves. DynaWM outperforms ablation variants in prediction loss and episode reward.
Real-world experiments validate successful traversal of continuous stairs with varying widths and heights, demonstrating smoother motion and better terrain adaptability.
Limitations
- Limited training terrain diversity: While five terrain types are included, continuous stair height/width ranges are limited (step height ≤0.18m), and more extreme stair configurations are not evaluated.
- World model prediction accuracy depends on terrain complexity: For highly dynamic or unstructured terrain, the residual network world model may not accurately predict future states, affecting teacher encoder dynamics-awareness quality.
- Limited hardware validation scale: Real-world experiments are conducted on a single bipedal-wheeled robot platform without systematic cross-platform generalization evaluation.
Conclusion and Outlook
DynaWM addresses dynamics representation weakening and dimensional collapse in teacher-student distillation through two core innovations. The world model as forward-dynamics regularizer compels the teacher encoder to retain geometric information essential for terrain dynamics modeling rather than merely maximizing immediate rewards. The momentum target encoder provides stable distillation targets through the BYOL paradigm, enabling robust student learning under non-stationary teacher supervision. PCA visualization quantitatively validates the encoder's hierarchical terrain height encoding capability.
On continuous stair traversal, DynaWM demonstrates superior terrain adaptability and motion smoothness, validated by real-world deployment. This framework provides an interpretable, verifiable methodology for terrain-aware representation learning in legged robots.
Golden insight: A good representation must not only "see" the terrain but "understand" how it will change—DynaWM uses a world model to teach the encoder to predict the future, and momentum targets to keep distillation stable in turbulent waters.
SOURCE LINKS
