🤖 AI

DynaWM: Dynamics-Aware Distillation with World Model and Momentum Targets for Smooth Locomotion over Continuous Stairs

This paper introduces DynaWM, a dynamics-aware distillation framework that combines a world model regularizer and momentum target encoding to enhance terrain representation and stabilize knowledge transfer, enabling bipedal-wheeled robots to achieve smooth and adaptive locomotion over continuous stairs.

Original authors: Haidong Hou, Zhangguo Yu, Hengbo Qi, Jianlin Zhang

Published 2026-06-24

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Haidong Hou, Zhangguo Yu, Hengbo Qi, Jianlin Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a robot that is part bicycle and part human: it has wheels for cruising on flat ground but also has legs to step over obstacles. This is a bipedal-wheeled robot. While these robots are great at handling flat floors or single steps, they often stumble when faced with a long, continuous staircase. They get shaky, lose their balance, or simply can't figure out how to climb smoothly.

The paper introduces a new method called DynaWM to teach these robots how to climb stairs like a pro. Here is how it works, explained through simple analogies.

The Problem: The "Black Box" Teacher

Currently, robots learn to move using a "Teacher-Student" system.

The Teacher: A super-smart AI that can see the stairs clearly (using cameras and sensors) and knows exactly how the robot's body should react. It's like a driving instructor who can see the whole road.
The Student: The robot's actual brain, which only feels its own body (like knowing its joints are bent) but cannot see the stairs. It has to guess what to do based on how the Teacher acts.

The Flaw: The current Teachers are too focused on the immediate reward (getting to the top of the step right now). They ignore the "big picture" of the terrain's shape. Also, because the Teacher changes its mind very quickly while learning, the Student gets confused and starts making simple, repetitive mistakes (like a student trying to copy a teacher who is constantly changing their handwriting).

The Solution: DynaWM

The authors built a new training system with two main upgrades to fix these issues.

1. The "Crystal Ball" (The World Model)

To make the Teacher smarter, they gave it a World Model. Think of this as a crystal ball or a flight simulator inside the Teacher's head.

How it works: Before the Teacher tells the Student what to do, it asks the crystal ball: "If I move my leg like this, what will the terrain look like in the next second?"
The Result: This forces the Teacher to pay attention to the actual shape and physics of the stairs (the "dynamics"), not just the immediate reward. It stops the Teacher from ignoring important details just because they don't give an instant point. It ensures the robot understands the geometry of the stairs, not just the goal.

2. The "Steady Hand" (Momentum Targets)

To stop the Student from getting confused by the Teacher's rapid changes, they introduced a Momentum Target.

The Analogy: Imagine trying to learn a dance routine from a teacher who is constantly changing the steps every second. You would never learn. Instead, DynaWM uses a "Steady Hand" version of the Teacher. This version is a slow-moving average of the Teacher's recent moves.
The Result: The Student learns from this calm, consistent version. Even if the real Teacher is frantic and changing fast, the Student gets a smooth, stable target to copy. This prevents the Student's brain from "collapsing" into a state where it stops learning complex patterns.

The Results: Smooth Climbing

The team tested this on a real robot (named JiaRan) and in simulations.

The Test: They threw the robot at various staircases, including some with uneven edges and heights it had never seen before.
The Outcome:
- Better Vision: The robot's internal "map" of the stairs became much clearer. Instead of a messy jumble of data, the robot organized the information by height, like a well-organized library.
- Smoother Motion: Compared to older methods, the robot climbed with much less shaking. It moved fluidly, like a human walking up stairs, rather than a robot jerking its way up.
- Success Rate: It successfully climbed continuous stairs up to 20cm high (about 8 inches) with a 100% success rate in real-world tests, whereas other methods often failed or fell.

In Summary

DynaWM is like giving a robot a driving instructor who not only knows the road but also simulates the future to understand the terrain better, while providing a calm, steady guide so the robot doesn't get overwhelmed. The result is a robot that can confidently and smoothly walk up long, tricky staircases without falling.

Technical Summary: DynaWM: Dynamics-Aware Distillation with World Model and Momentum Targets for Smooth Locomotion over Continuous Stairs

Problem Statement
While bipedal-wheeled robots have demonstrated success in traversing slopes and single-step obstacles using model predictive control and deep reinforcement learning, navigating long, continuous staircases remains a significant challenge. Current solutions relying on teacher-student distillation frameworks face three primary limitations:

Weakened Dynamics-Aware Representations: Teacher encoders optimized solely through policy gradients tend to encode only terrain features directly relevant to immediate rewards, disregarding broader geometric information essential for ground dynamics modeling.
Dimensional Collapse: In concurrent teacher-student training, rapid updates to the teacher's representations can cause the student encoder to collapse into trivial representations (fitting only mean predictions) rather than capturing essential terrain features.
Opacity: The latent space in existing frameworks acts as a "black box," making it difficult to verify the true modeling capability for terrain height or ensure interpretability during training.

These issues result in a lack of motion smoothness and terrain adaptability when robots face continuous stairs, hindering reliable sim-to-real deployment.

Methodology: DynaWM
The authors propose DynaWM (Dynamics-aware distillation with World model and Momentum targets), a representation learning framework designed to enhance terrain encoding and stabilize knowledge transfer. The framework consists of three interconnected modules:

World Model as a Regularizer:
To enforce forward-dynamics awareness, a world model is introduced as a regularizer for the teacher encoder. The teacher encoder ( $E_{\theta_t}$ ) generates latent variables ( $z_t^t$ ) which, combined with privileged observations, are fed into a world model ( $W_\omega$ ) to predict the next privileged state ( $\hat{s}_{t+1}$ ).
- Mechanism: A prediction loss ( $L_{pred}$ ) is minimized to train the world model and propagate gradients back to the teacher encoder.
- Effect: This compels the encoder to retain information essential for forward dynamics (e.g., terrain geometry and friction) rather than discarding it as noise irrelevant to immediate rewards. The world model is implemented as a deep residual network with Dense-LayerNorm-Swish blocks to capture long-range temporal dependencies.
Momentum Target Encoder for Stable Distillation:
To address the non-stationarity of the teacher encoder during concurrent training, the framework employs a momentum target encoder inspired by BYOL (Bootstrap Your Own Latent).
- Mechanism: Instead of the student regressing directly to the rapidly changing online teacher, it learns from a target encoder ( $E_{\hat{\theta}_s}$ ) that maintains an exponential moving average (EMA) of the student's own weights. A predictor head is added to the online student to prevent collapse.
- Effect: This provides consistent, stable distillation targets, preventing dimensional collapse and allowing the student to learn robust dynamics-aware representations from proprioceptive history alone.
Reward Design and Training:
The reward function is augmented with specific constraints (joint range, wheelbase distance, base position bounds, and leg balance) to mitigate issues like excessive hip reliance and body oscillation observed in stair climbing. Training utilizes a Concurrent Teacher-Student (CTS) paradigm with PPO, incorporating domain randomization across terrain types (flat, slopes, steps, uneven) and physical parameters (mass, friction, PD gains) to bridge the sim-to-real gap.

Key Contributions

World Model Regularization: The introduction of a world model to enforce forward-dynamics awareness, ensuring learned representations capture terrain geometry and dynamics rather than solely maximizing rewards.
Stable Distillation via Momentum Targets: The use of a momentum target encoder to provide stable distillation targets, effectively preventing dimensional collapse caused by the teacher's non-stationary updates.
Transparent Assessment: The application of Principal Component Analysis (PCA) and quantitative metrics (PC-Correlation, Linear Probing $R^2$ , CCA) to verify that the encoders hierarchically capture terrain height, offering a transparent view of representation quality.

Experimental Results
The method was evaluated in simulation (IsaacGym) and on a real-world JiaRan bipedal-wheeled robot.

Representation Quality: PCA visualizations and quantitative metrics showed that the proposed method achieved a PC-Correlation of 0.96 for terrain height in the teacher encoder and 0.87 in the student encoder, significantly outperforming ablations without the world model or momentum targets (which suffered from entangled representations or collapse).
Simulation Performance: On diverse continuous stairs (varying widths and heights up to 20 cm), DynaWM achieved success rates up to 100% for 26cm wide steps and 89.3% for 20cm high wide steps, outperforming baselines like CTS and RMA.
Motion Smoothness: The method demonstrated superior motion smoothness (lower jerk and base acceleration) and comparable energy consumption compared to CTS.
Real-World Deployment: In real-world experiments on unseen terrains, including irregular stairs with protrusions and continuous steps exceeding training heights, the robot achieved a 100% success rate (5/5 trials) with significantly smoother motion than the CTS baseline.

Significance and Claims
The paper claims that DynaWM successfully bridges the gap between privileged learning and deployable control for complex terrain. By ensuring the encoder captures comprehensive terrain geometry through dynamics-aware regularization and stabilizing the distillation process, the method enables bipedal-wheeled robots to traverse lengthy, continuous staircases with high adaptability and smoothness. The authors assert that this approach overcomes the limitations of current teacher-student frameworks, which often fail to maintain motion smoothness and terrain adaptability in continuous stair scenarios. The work demonstrates that explicit modeling of forward dynamics and stable target generation are critical for robust sim-to-real transfer in legged-wheeled locomotion.