Research map / Reinforcement learning / Visual and representation RL
Discrete Codebook World Models for Continuous Control
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DCWM, a self-supervised world model whose latent states are discrete FSQ codebook codes trained with cross-entropy and stochastic dynamics, combined with MPPI planning in DC-MPC, competitive with TD-MPC2 and DreamerV3 on DMControl, Meta-World and MyoSuite.
Problem
World models for state-based continuous control: discrete latent spaces (DreamerV3) underperform continuous latent spaces (TD-MPC2), and it is unclear whether discrete latents trained with classification help.
Approach
Represent latent states as finite scalar quantization (FSQ) codes from a fixed codebook, model stochastic latent transitions as a categorical classifier trained with cross-entropy using straight-through Gumbel-softmax multi-step rollouts, and plan with MPPI using expected codes.
Key results
- DC-MPC outperforms DreamerV3 on Meta-World manipulation tasks and performs strongly on high-dimensional DMControl Dog and Humanoid tasks, while generally matching TD-MPC2 at 1M environment steps.
- Discrete latent spaces (FSQ with cross-entropy) show significantly better sample efficiency than continuous latent spaces trained with MSE regression in the latent-space comparison.
- Replacing the codebook encoding with label or one-hot encodings for the reward, value and policy networks causes flat-lining training (label or one-hot fail to learn), supporting the codebook encoding.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a discrete codebook latent space instead of a continuous latent space (ablation) | DC-MPC world model · DMControl | sample efficiency | Helped |
| Quantize the latent space with FSQ (ablation) | DC-MPC world model · DMControl | sample efficiency | Helped |
| Use stochastic categorical dynamics with classification instead of deterministic discrete dynamics with MSE (ablation) | DC-MPC world model · continuous control tasks | performance | Helped |
| Use stochastic Gaussian or Gaussian-mixture dynamics instead of deterministic continuous dynamics (ablation) | continuous-latent world models · DMControl and Meta-World | sample efficiency | Mixed |
| Replace codebook encodings with label encodings for reward, critic, and prior-policy inputs (ablation) | DC-MPC · Humanoid Walk | sample efficiency | Hurt |
| Replace codebook encodings with one-hot encodings for reward, critic, and prior-policy inputs (ablation) | DC-MPC · continuous control tasks | sample efficiency | No clear effect |
| Use a one-hot encoding instead of a codebook encoding in DreamerV3 (ablation) | DreamerV3 · Reacher Hard and Walker Walk | performance | No clear effect |
| Replace DreamerV3's one-hot encoding with the DC-MPC codebook encoding (ablation) | DreamerV3 · continuous control tasks | performance | No clear effect |
| Increase codebook size (ablation) | DC-MPC · continuous control tasks | learning speed | Mixed |
| Use a smaller latent dimension (ablation) | DC-MPC · difficult environments, including Humanoid Walk | performance | Hurt |
| Replace FSQ with vector quantization (ablation) | DC-MPC · Walker Walk, Dog Run, Humanoid Walk, and Reacher Hard | performance | Mixed |
| Use a 5-critic REDQ ensemble instead of standard double Q with 2 critics (ablation) | DC-MPC · Dog Run and Humanoid Walk | performance | Helped |
Surprising findings
- Label and one-hot encodings in the self-supervised world model fail to learn the task at all when used as the sole encoding, while the codebook encoding works.
- Stochastic continuous latent spaces (Gaussian, GMM) sometimes help on DMControl but drastically underperform on Meta-World.
Limitations
- Hyperparameters are task-specific (noise schedules and N-step returns tuned per task).
- Does not model epistemic uncertainty or include principled exploration.
- Observation reconstruction is not used; results with reconstruction were worse, and DreamerV3 underperforms in hard tasks.
Experimental setup
Datasets: DeepMind Control Suite Meta-World MyoSuite
Benchmarks: DeepMind Control Suite Meta-World MyoSuite
Metrics: normalized return success rate median IQM mean optimality gap
Baselines: TD-MPC2 DreamerV3 TD-MPC SAC
Builds on: TD-MPC2 DreamerV3 FSQ REDQ TD3
Closest papers by meaning
- TD-MPC2: Scalable, Robust World Models for Continuous Control (ICLR 2024 · similarity 0.87)
- PWM: Policy Learning with Multi-Task World Models (ICLR 2025 · similarity 0.87)
- Overcoming Slow Decision Frequencies in Continuous Control: Model-Based Sequence Reinforcement Learning for Model-Free Control (ICLR 2025 · similarity 0.87)
- Simplifying Latent Dynamics with Softly State-Invariant World Models (NeurIPS 2024 · similarity 0.86)
- Reinforcement World Model Learning for LLM-based Agents (NeurIPS 2026 · similarity 0.86)
- Masked Diffusion Language Models are Strong and Steerable Text-Based World Models for Agentic RL (NeurIPS 2026 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.