Pub-AI.AI in Science

Latent Dynamics · 2026-10-20

Dynamics-Aligned Representation Learning under Induced Partial Observability

SM Mazharul Islam, Manfred Huber

arXiv:2610.25511PDF

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

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

preprint

Dynamics-Aligned Representation Learning under Induced Partial Observability

SM Mazharul Islam, Manfred Huber

Preprint posted: 2026-10-18 · Original record ↗

arxiv
2610.25511
Explicit source versions (1)
  • v1 · 2026-10-18
  1. Preprint available

    Source date: 2026-10-18 · First detected: 2026-10-25 UTC · Historical baseline

    Source evidence (1)
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.