Latent Dynamics · 2026-10-20
Dynamics-Aligned Representation Learning under Induced Partial Observability
SM Mazharul Islam, Manfred Huber
AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.
TL;DR
The authors propose stochastic spatial gating for latent dynamics models in model-based RL. During training, they mask pixels using a binary mask sampled from a Beta–Bernoulli random field (mask probability resampled each timestep) to induce controlled partial observability. They minimally integrate the method into RSSM by concatenating the gated observation and the mask to the encoder and adapting the reconstruction loss to only surviving pixels. Evaluation uses MPC-CEM open-loop rollouts on DeepMind Control Suite tasks.
Why it matters
Latent dynamics models are often trained with observation reconstruction losses that can bias representations toward visual fidelity, undertrain the prior, and produce unreliable long-horizon open-loop predictions for planning. By inducing partial observability during training (without major pipeline changes) and aligning the reconstruction loss with surviving pixels, the authors aim to prevent posterior dominance/collapse and improve dynamics-aligned representations for planning and data efficiency.
Method
- Induce induced partial observability during training by sampling a spatial binary mask G(t) from a Beta–Bernoulli random field; resample the gate probability at each timestep and apply the mask to raw pixel observations.
- Minimal RSSM integration: concatenate the gated observation (G(t) ⊙ o_t) and the binary mask G(t) as a channel-wise augmented input to the encoder.
- Use a gating-aware observation reconstruction loss that restricts reconstruction to observed pixels, and stabilize prior learning with KL balancing (set α=0.8) and free bits.
Abstract (from arXiv)
Model-based reinforcement learning (MBRL) has enabled sample-efficient learning of continuous control policies from high-dimensional observations through learned latent dynamics models. These models are typically trained with observation reconstruction objectives, which provide strong supervision but can bias representations toward visual fidelity leading to an undertrained prior and unreliable long-horizon predictions. We propose a stochastic gating mechanism that induces controlled partial observability during training by sampling a spatial binary mask from a Beta--Bernoulli random field and applying it to the raw observation. The gate probability is resampled at each time step, exposing the encoder to conditions ranging from near-total occlusion to near-full visibility. Integration of our method requires minimal architectural changes from existing dynamics models: the gated observation and the binary mask are channel-wise concatenated and passed to the encoder as a single augmented input, with no modifications to the broader training pipeline. We evaluate the proposed approach on a set of pixel-based continuous control tasks selected from the DeepMind Control Suite. Using model predictive control for evaluation, which relies solely on open-loop model rollouts, we show that stochastic gating improves planning performance and data efficiency across tasks. Further analysis demonstrates that performance is robust to the spatial scale of the mask, and a balanced observability during training outperforms configurations biased toward either sparser or denser observations.
Related papers
- Learning Transferable Policies from Action-free Time Series Through Dynamical Embeddings
- VIGOR: Zero-Shot Visual Generalization via Latent-Space Consistency in Model-Based Reinforcement Learning
- Directed Temporal Representations for Offline Visual Control
- World Models with Configurable Perception
- Preserving Unstable Modes Through Inverse Dynamics in JEPA World Models
Publication record
Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.
First public release: 2026-10-18 · Indexed by Pub-AI: 2026-10-20 UTC
Dynamics-Aligned Representation Learning under Induced Partial Observability
Preprint posted: 2026-10-18 · Original record ↗
- arxiv
- 2610.25511
Explicit source versions (1)
- v1 · 2026-10-18
- Preprint available
Source date: 2026-10-18 · First detected: 2026-10-25 UTC · Historical baseline
Source evidence (1)
- Official evidence ↗ · atom:published (official metadata)
Abstract · arxiv
Model-based reinforcement learning (MBRL) has enabled sample-efficient learning of continuous control policies from high-dimensional observations through learned latent dynamics models. These models are typically trained with observation reconstruction objectives, which provide strong supervision but can bias representations toward visual fidelity leading to an undertrained prior and unreliable long-horizon predictions. We propose a stochastic gating mechanism that induces controlled partial observability during training by sampling a spatial binary mask from a Beta--Bernoulli random field and applying it to the raw observation. The gate probability is resampled at each time step, exposing the encoder to conditions ranging from near-total occlusion to near-full visibility. Integration of our method requires minimal architectural changes from existing dynamics models: the gated observation and the binary mask are channel-wise concatenated and passed to the encoder as a single augmented input, with no modifications to the broader training pipeline. We evaluate the proposed approach on a set of pixel-based continuous control tasks selected from the DeepMind Control Suite. Using model predictive control for evaluation, which relies solely on open-loop model rollouts, we show that stochastic gating improves planning performance and data efficiency across tasks. Further analysis demonstrates that performance is robust to the spatial scale of the mask, and a balanced observability during training outperforms configurations biased toward either sparser or denser observations.