Research map / Multi-agent systems and games / Human-AI and goal-conditioned
Simplifying Latent Dynamics with Softly State-Invariant World Models
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 PLSM, a world-model regularizer that penalizes the norm of a query representation of latent-action pairs so that action-induced latent dynamics depend minimally on the state (softly state-invariant), and evaluates it with contrastive world models, TD-MPC/RePo planning, and SPR model-free RL.
Problem
World models often do not represent the effects of actions in systematic, predictable ways in latent space, which hurts planning, generalization, and robustness.
Approach
Minimize the mutual information between the latent state and the predicted latent change conditioned on the action, implemented by routing dynamics through a query network whose output h_t is L2-penalized (a variational KL-to-Gaussian bound), so the dynamics become softly state-invariant.
Key results
- On Atari 100k, SPR+PLSM raises mean human normalized score from 61.5% to 67.1% across 26 games.
- PLSM improves contrastive world model long-horizon latent prediction accuracy in five of six environments; with slots at horizon 10 on cubes 9, PLSM reaches 99.05% vs 97.45% for C-SWM.
- In high noise (sigma=0.2) on dSprites, PLSM dynamics predict almost 60% of transitions after 10 steps versus under 50% for unconstrained models.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add PLSM dynamics to SPR | SPR · Atari 26 games | mean human-normalized score: 0.616 -> 0.671 | Helped |
| Add PLSM information regularization to slot-based C-SWM | C-SWM · Cubes 9, prediction horizon t=10 | accuracy: 97.45%±1.3 -> 99.05% ±0.18 | Helped |
| Add PLSM dynamics regularization | TD-MPC · DeepMind Control Suite continuous-control tasks | return | Helped |
| Add PLSM dynamics regularization | TD-MPC · quadruped-run | return | No clear effect |
| Add PLSM dynamics regularization | RePo · Distracting Control Suite, especially hopper-stand, walker-run, and finger-spin | return | Helped |
| Add PLSM dynamics regularization | Contrastive World Model · Six latent-prediction environments | Hits at Rank 1 accuracy | Helped |
| Add PLSM dynamics regularization | Contrastive World Model · Pong | Hits at Rank 1 accuracy | No clear effect |
| Add PLSM dynamics regularization | Contrastive World Model · Cubes and shapes with fewer objects than training data | accuracy | Helped |
| Add PLSM dynamics regularization | Contrastive World Model · dSprite with high Gaussian noise | transition prediction accuracy at 10 steps: almost 60% ... less than 50% | Helped |
| Remove the query representation and regularize the latent state instead (ablation) | PLSM · Three tested datasets | performance | Hurt |
| Replace the full query representation with a top-15-features bottleneck (ablation) | PLSM · Tested datasets | performance | No clear effect |
| Add weight decay to the dynamics MLP (ablation) | PLSM · Tested datasets | performance | No clear effect |
Surprising findings
- Removing the query representation and regularizing the latent state instead substantially hurts performance.
- PLSM reduces performance on some Atari games (e.g., BankHeist, ChopperCommand), and Pong showed no advantage due to non-controllable elements.
Limitations
- Assumes actions have predictable, deterministic effects; uncontrollable aspects may be ignored.
- Regularization coefficient beta must be tuned per environment.
- Performance drops on some Atari games such as BankHeist and ChopperCommand.
Experimental setup
Datasets: DeepMind Control Suite Distracting Control Suite Atari 100k Pong Space Invaders cubes and shapes grid worlds dSprites MOVi-E DAVIS 2017
Benchmarks: DeepMind Control Suite Distracting Control Suite Atari 100k cubes and shapes dSprites MOVi-E
Metrics: average return human normalized score Hits at Rank 1 accuracy IQM optimality gap
Baselines: TD-MPC RePo SPR CURL DrQ CWM C-SWM
Builds on: TD-MPC Contrastive World Models SPR RePo Dreamer
Closest papers by meaning
- HaM-World: Soft-Hamiltonian World Models with Selective Memory for Planning (NeurIPS 2026 · similarity 0.89)
- Beyond Prediction: Steering VLM Agents with Retrospective World Modeling (NeurIPS 2026 · similarity 0.89)
- Reinforcement Learning Under Latent Dynamics: Toward Statistical and Algorithmic Modularity (NeurIPS 2024 · similarity 0.89)
- Revisiting Multi-Agent World Modeling from a Diffusion-Inspired Perspective (NeurIPS 2025 · similarity 0.88)
- General agents need world models (ICML 2025 · similarity 0.88)
- World Action Verifier: Self-Improving World Models via Forward-Inverse Asymmetry (NeurIPS 2026 · similarity 0.88)
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.