Coupled Distributional Random Expert Distillation for World Model Online Imitation Learning

TL;DR

The paper tackles instability in world-model-based imitation learning by replacing adversarial rewards with a joint, latent-space density estimation via Random Network Distillation (RND). The proposed Coupled Distributional Random Expert Distillation (CDRED) jointly estimates expert and behavioral distributions using a shared ensemble of random targets, and integrates this with decoder-free world models and MPPI planning. Key contributions include a theoretically grounded unbiased estimator for online data-frequency, a two-predictor reward model with a dual-target structure, and empirical demonstrations of stability and expert-level performance across DMControl, Meta-World, and ManiSkill2. The work delivers a robust online imitation framework with strong planning capabilities, reducing issues like overly strong discriminators and long-term instability seen in adversarial approaches, while enabling reliable control in both locomotion and manipulation tasks.

Abstract

Imitation Learning (IL) has achieved remarkable success across various domains, including robotics, autonomous driving, and healthcare, by enabling agents to learn complex behaviors from expert demonstrations. However, existing IL methods often face instability challenges, particularly when relying on adversarial reward or value formulations in world model frameworks. In this work, we propose a novel approach to online imitation learning that addresses these limitations through a reward model based on random network distillation (RND) for density estimation. Our reward model is built on the joint estimation of expert and behavioral distributions within the latent space of the world model. We evaluate our method across diverse benchmarks, including DMControl, Meta-World, and ManiSkill2, showcasing its ability to deliver stable performance and achieve expert-level results in both locomotion and manipulation tasks. Our approach demonstrates improved stability over adversarial methods while maintaining expert-level performance.

Figures (15)

• Figure 1: Coupled Distributional Random Expert Distillation We present the architecture of our CDRED reward model. During training, the behavioral and expert predictors are trained using latent representations encoded from observations and actions sampled from the behavioral and expert buffers. The dotted blue lines indicate the gradient backpropagation paths. During inference, rewards are estimated by the outputs of the behavioral and expert predictors, along with the mean and second-order moments of the target network's output, for an unseen latent state-action pair.
• Figure 2: Intuitive Illustration for Coupled Distribution Estimation When the state-action distribution of the initial policy differs significantly from that of the expert distribution, the initial rewards tend to approach zero. This often leads to a slower or even unsuccessful learning process. By estimating the behavioral distribution in conjunction with the expert distribution, we can effectively model the rewards to guide the behavioral distribution closer to that of the expert.
• Figure 3: Advantages of Coupled Density Estimation We demonstrate the empirical performance boost of our coupled density estimation in terms of leveraging random network distillations for reward modeling based on state-action distribution estimation. With coupled estimation, we observe faster convergence to optimal in many simple cases (Left) and better performance in complex tasks (Right).
• Figure 4: Meta-World Results We evaluate our CDRED method (red lines) on 6 tasks in Meta-World environments. We show stabler performance on these tasks, outperforming the baselines. IQ-MPC (orange lines) suffers from overly powerful discriminator problem mentioned in Section \ref{['sec:drawbacks-iqmpc']}. We conduct the experiments on 3 random seeds.
• Figure 5: DMControl Results We evaluate our CDRED method (red lines) on 6 tasks in DMControl environments. Our approach achieves results comparable to IQ-MPC (orange lines) in Hopper Hop, Walker Run, and Humanoid Walk, while demonstrating greater stability across the remaining tasks. We conduct the experiments on 3 random seeds.

Theorems & Definitions (2)

• Lemma 1: Unbiased Estimator
• proof