Shaping Belief States with Generative Environment Models for RL

Motivation

• The target representation should capture temporally coherent global factors, such as an environment map, despite rapidly changing first-person observations.The motivation is task-independent structure that can support many downstream tasks.Source: §1 · p. 1
• Long-term consistency is a central challenge because models may predict short-term observations accurately while failing to maintain coherent predictions in partially observed environments.The paper attributes this problem to conditioning and weak objectives rather than model capacity alone.Source: §1 · p. 2

Contributions

• The paper proposes learning expressive generative models from first-person views and actions to capture long-term consistency in rich 3D environments.It analyzes belief-state architectures, overshooting, memory, and shared representations between the model and agent.Source: §1 · p. 2
• Sharing the generative model’s belief-state with the agent substantially increases data-efficiency relative to strong model-free baselines without significantly affecting training speed.The contribution is evaluated across a variety of reinforcement-learning tasks.Source: §1 · p. 2

2 Methods

• The architecture combines an online recurrent agent state with a SimCore recurrent network that predicts future frame distributions from actions and a starting belief-state.The agent uses the belief-state for policy and value, while SimCore simulates forward and conditions the generative frame model.Source: §2 · p. 2, Figure 1 · p. 3

2.1 Frame generative models and the problem of conditioning.

• Expressive generative models can ignore conditioning, making belief-state gradients approximately zero; modified ConvDRAW with GECO learns stable states while retaining sample quality.The paper uses the resulting model to simulate multiple steps into the future.Source: §2.1 · pp. 2–3
• The training objective samples future points along trajectories and uses an overshoot length of typically 4-32 steps while limiting evaluation points and unrolls for computational efficiency.Table 1 identifies typically 2 future evaluation points and 6 random trajectory locations.Source: Table 1 · p. 3

2.2 Why is next-step prediction not sufficient?

• Next-step prediction weakens dependence on the belief-state because immediate observations can be inferred from recent history, whereas distant predictions require global environmental structure.This motivates overshooting rather than relying only on one-step prediction.Source: §2.2 · pp. 3–4
• Overshooting increases target uncertainty, leading the authors to hypothesize that generative predictors benefit more from longer horizons than deterministic predictors.The paper notes that prior overshooting evidence emphasized performance or long-term prediction rather than which environment aspects are retained.Source: §2.2 · p. 4

2.3 Belief-states, Localization and Mapping

• The method decodes position, orientation, and top-down layout from the belief-state without restricting the representation to 2D maps or interfering with agent training.Different decoders and memory reads are used for LSTM, slot-based, and Kanerva-memory states.Source: §2.3 · p. 4

Model-based RL

• Model-based RL has substantial prior support, but demonstrating benefits in visually rich, dynamic 3D environments remains challenging without privileged information.Prior systems have used camera locations, depth, or simulator state variables.Source: §3 · p. 4

Auxiliary learning

• The paper treats its generative model as an auxiliary representation-shaping loss for a model-free RL agent rather than using the model directly for planning.This connects the approach to pixel control, CPC, frame reconstruction, next-frame prediction, and successor representations.Source: §3 · p. 4

Distinction from prior work

• Unlike closely related comparisons of action-conditional CPC and deterministic next-step prediction, this work analyzes map decoding, expressive generative simulation, overshooting, and RL performance.It shares a belief-state between the generative model and policy network.Source: §3 · pp. 4–5

4 Experiments

• The experiments vary unsupervised loss, overshoot length, and belief-state memory across DeepMind Lab, Random City, voxel building, and random-terrain environments.Training uses IMPALA with the model jointly trained through a shared belief network, reducing running speed by about 20−40%.Source: §4 · p. 5

4.1 Random City

• Belief-state reconstructions preserve previously observed buildings outside the current view and support forward simulation in Random City, despite training only from first-person views.The experiment uses fixed handcrafted policies rather than reinforcement learning.Source: §4.1 · p. 5, Figure 2 · p. 5
• In Random City, generative prediction with Kanerva memory provides the best map decoding, while contrastive loss provides the best position decoding.Increasing overshoot improves position and map decoding up to a certain length, with generative prediction especially sensitive for map accuracy.Source: §4.1 · p. 6, Figure 3 · p. 6
• Contrastive loss localizes accurately but maps poorly, whereas Kanerva memory outperforms pure LSTM and slot-based memory on map decoding.The authors caution that slot-based memory may retain information that is difficult for the map-decoding analysis to extract.Source: §4.1 · p. 6

4.2 DeepMind Lab

• Generative SimCore improves performance for both LSTM and RMA agents in DeepMind Lab, while Kanerva memory is unstable in the RL setting.Longer overshoots improve map reconstruction but do not necessarily improve RL performance.Source: §4.2 · p. 7

4.3 Voxel environment

• In the voxel environment, SimCore significantly improves performance across all four levels and makes performance less sensitive to Adam hyperparameters and unroll length.The agent learns building behaviors such as towers and stairs, and can simulate movement, construction, and block picking.Source: §4.3 · p. 8, Figure 7 · p. 8
• In a harder procedurally generated terrain, an LSTM retains an approximate map, but information not observed for a while gradually fades.The authors identify scaling these experiments as future work.Source: §4.3 · p. 8

5 Discussion

• Expressive generative models combined with overshooting form stable belief states in 3D environments from first-person views, while sharing those states improves RL data-efficiency.The representation captures environmental structure with little prior knowledge, but further elements require investigation.Source: §5 · p. 8
• Jointly training the belief state with the agent makes belief formation or map decoding harder.The paper suggests policy gradients, policy changes, or changing representations may contribute.Source: §5 · p. 9
• Future work includes scaling the system through improved training or memory architectures and using the model for planning.Source: §5 · p. 9

Appendix A Modified ConvDRAW

• ConvDRAW uses recurrent convolutional encoder and decoder networks with latent variables at each step to form autoregressive latent distributions.It is an instance of a deep variational auto-encoder.Source: Appendix A · p. 14
• The decoder canvas accumulates reconstruction information into a distribution over inputs p(x|z).Its operations and states are convolutional, with spatial and feature dimensions.Source: Appendix A · p. 14
• The model replaces LSTM recurrence with tanh-sigmoid nonlinearities, uses a separate prior network for conditioning, and adaptively balances input and latent losses.These changes reduce operations, improve conditioning, and target reconstruction accuracy.Source: Appendix A · p. 14
• Figure 8 diagrams ConvDRAW’s likelihood computation.Source: Figure 8 · p. 14
• The implementation performs eight repeated operations on 64 × 64 inputs using convolutional encoder, decoder, recurrent, and canvas components.The reconstruction error is the difference between the current canvas and reconstruction.Source: Appendix A · p. 14

Appendix B Memory Architectures

• Classic LSTMs store rapidly changing information in activations because their weights update slowly through back-propagation.Source: Appendix B · p. 14
• Slot-based memory allocates a new slot at every time step and retrieves stored vectors through attention.It avoids forgetting and does not require learning how to write, but allocates slots even when nothing new occurs.Source: Appendix B · pp. 14–15
• Kanerva memory treats a global latent variable as memory and updates it by exact, efficient posterior inference in a linear Gaussian model.Source: Appendix B · p. 15

Appendix C Integrating Kanerva Machines with the RNNs

• The Kanerva module is trained end-to-end by back-propagation with the rest of the model without an auxiliary loss.It functions as external memory for the SimCore.Source: Appendix C · p. 15
• At each step, the integrated Kanerva machine reads from and writes to memory using words computed from recurrent state, action, and observation embedding.The read word queries memory, while the write word updates its posterior.Source: Appendix C · p. 15
• Figure 9 illustrates a dynamic Kanerva machine integrated with an RNN.Source: Figure 9 · p. 15

Appendix D Reinforcement learning framework

• At each discrete time step, the agent observes a frame, produces a discrete action, receives a reward, and updates from its belief state.Source: Appendix D · p. 15
• The objective maximizes the expected discounted sum of rewards conditioned on the belief state, with policy entropy regularization.Source: Appendix D · p. 16
• IMPALA training uses parallel actors that collect trajectories and one learner that batches them, unrolls recurrent cores, and computes losses and gradients.Each experience item is used once and is only slightly off-policy.Source: Appendix D · p. 16

• Curves are first smoothed with an exponential moving window of length 10.Source: Appendix E · p. 16
• Smoothed curves are then resampled at 128 uniformly spaced x-axis points using linear interpolation.Source: Appendix E · p. 16
• Error bars represent the 90% confidence region.Source: Appendix E · p. 16

• Table 2 reports the hyper-parameters used in the experiments.Source: Appendix F · p. 16, Table 2 · p. 16
• Table 3 reports the parameters for the slot-based RMA memory.Source: Appendix F · p. 16, Table 3 · p. 16
• Table 4 reports the parameters for Kanerva memory.Source: Appendix F · p. 16
• Each reported experiment was repeated at least 3 times with different random seeds.Source: Table 2 · p. 16

Appendix G GECO

• GECO constrains reconstruction error with threshold κ through a Lagrangian min-max formulation rather than direct ELBO maximization.The formulation uses qφ(z|x) as an approximation to the posterior pθ(z|x).Source: Appendix G · pp. 16–17
• κ has an empirically optimal range for map reconstruction, with κ ≈1e-3 producing the best result in Figure 10.Source: Figure 10 · p. 17, Appendix G · p. 17

Appendix H Map decoding on Random City environment

• Additional map decoding samples are shown for each model in Figure 11.Source: Appendix H · p. 17

Appendix I Procedurally generated terrain

• A 64 × 64 procedurally generated terrain was created to analyze the belief state in a more naturalistic environment.Source: Appendix I · p. 17
• Map decoding MSE decreases along a terrain trajectory, indicating accumulation and retention of evidence about environmental layout.Source: Appendix I · p. 17

• The agent learns harder levels involving longer staircases and taller structures for elevated food locations.Source: Appendix J · p. 17
• Model rollouts are simulated deterministically in latent space and rendered by sampling from ConvDRAW.Samples should match observed frames when the simulation knows the scene and differ while remaining knowledge-consistent under uncertainty.Source: Appendix J · p. 18
• In the terrain environment, reconstructed maps persist for some time but fade gradually when the agent does not revisit regions.Source: Figure 12 · p. 19
• Terrain map MSE decreases along a trajectory, showing accumulation and retention of layout evidence.Source: Figure 13 · p. 20
• Building-level samples depict stair construction, reaching high food, and constructing a tower to climb to a platform.Source: Figure 14 · p. 21
• Additional samples are provided for the terrain environment.Source: Figure 15 · p. 22