Learning to Predict Without Looking Ahead: World Models Without Forward Prediction

Abstract: Much of model-based reinforcement learning involves learning a model of an agent's world, and training an agent to leverage this model to perform a task more efficiently. While these models are demonstrably useful for agents, every naturally occurring model of the world of which we are aware---e.g., a brain---arose as the byproduct of competing evolutionary pressures for survival, not minimization of a supervised forward-predictive loss via gradient descent. That useful models can arise out of the messy and slow optimization process of evolution suggests that forward-predictive modeling can arise as a side-effect of optimization under the right circumstances. Crucially, this optimization process need not explicitly be a forward-predictive loss. In this work, we introduce a modification to traditional reinforcement learning which we call observational dropout, whereby we limit the agents ability to observe the real environment at each timestep. In doing so, we can coerce an agent into learning a world model to fill in the observation gaps during reinforcement learning. We show that the emerged world model, while not explicitly trained to predict the future, can help the agent learn key skills required to perform well in its environment. Videos of our results available at https://learningtopredict.github.io/

• The paper introduces observational dropout to implicitly train internal world models, bypassing explicit forward-predictive losses.
• The method modifies standard agent interactions by stochastically replacing true observations with null inputs to trigger reliance on hidden states.
• Experimental results show improved sample efficiency and robustness in environments like CarRacing and ViZDoom under intermittent sensory data.

This work investigates an alternative mechanism for inducing the learning of internal world models within reinforcement learning agents, deviating from the common practice of employing explicit forward-predictive losses. The core idea, termed "observational dropout," proposes that by intermittently restricting an agent's access to environmental observations, the agent can be implicitly incentivized to develop and rely upon an internal representation that captures environmental dynamics, effectively serving as a world model. This approach draws inspiration from biological systems where complex predictive capabilities emerge without direct supervision on future states but rather as a consequence of evolutionary pressures favouring robust behaviour in partially observable or noisy environments.

Observational Dropout Mechanism

The central mechanism introduced is observational dropout, a modification to the standard agent-environment interaction loop. In a typical Partially Observable Markov Decision Process (POMDP) setting, an agent at timestep $t$ receives an observation $o_t$, maintains an internal hidden state $h_t$, takes an action $a_t$ based on $o_t$ and $h_t$, transitions to a new hidden state $h_{t+1}$, and receives a reward $r_t$. The environment then transitions to a new state $s_{t+1}$ and emits observation $o_{t+1}$.

Observational dropout modifies this process by introducing a probability $p$ at each timestep $t$. With probability $p$, the agent does not receive the true observation $o_t$ from the environment. Instead, it receives a null or zero observation $\emptyset$. When this occurs, the agent must rely solely on its internal hidden state $h_t$ (updated from $h_{t-1}$ and the previous action $a_{t-1}$) and the null observation to select its next action $a_t$ and update its internal state to $h_{t+1}$. With probability $1-p$, the agent receives the true observation $o_t$ and proceeds as usual.

Formally, let $f$ be the recurrent state update function of the agent (e.g., an LSTM or GRU cell) and $\pi$ be the policy function. The standard update is:

$h_{t+1} = f(h_t, o_t, a_{t-1})$

$a_t = \pi(h_{t+1}, o_t)$

With observational dropout, the observation $\hat{o}_t$ used by the agent is determined as:

$$\hat{o}_t = \begin{cases} o_t & \text{with probability } 1-p \ \emptyset & \text{with probability } p \end{cases}$$

The agent's update rule then becomes:

$h_{t+1} = f(h_t, \hat{o}_t, a_{t-1})$

$a_t = \pi(h_{t+1}, \hat{o}_t)$

The critical aspect is that the agent is still trained using a standard reinforcement learning objective (e.g., maximizing expected cumulative reward via algorithms like A2C or PPO). The observational dropout is not part of the loss function itself but rather a modification of the data stream provided to the agent during rollouts. The hypothesis is that to maintain performance under this stochastic observation regime, the agent's recurrent state $h_t$ must learn to encode information predictive of the environment's state, effectively compensating for the missing observations. This internal state, therefore, emerges as an implicit world model.

Implementation Architecture

Implementing observational dropout requires an agent architecture capable of maintaining and utilizing an internal state across timesteps. A common choice is a recurrent neural network (RNN), typically an LSTM or GRU, integrated into the policy network.

A typical architecture consists of:

1. Observation Encoder: A network (e.g., CNN for visual input, MLP for vector input) that processes the raw observation $\hat{o}_t$ into a feature vector $e_t$. If $\hat{o}_t = \emptyset$, this encoder might receive a zero vector or a special token, producing a corresponding null embedding $e_t = e_{\emptyset}$.
2. Recurrent Core: An RNN (e.g., LSTM) that updates the hidden state $h_t$ based on the previous hidden state $h_{t-1}$, the previous action $a_{t-1}$, and the current encoded observation $e_t$: $h_t = \text{RNN}(h_{t-1}, [e_t, a_{t-1}])$. This hidden state $h_t$ is hypothesized to encapsulate the learned world model.
3. Policy Head: An MLP that takes the current hidden state $h_t$ as input and outputs the action distribution $\pi(a_t | h_t)$.
4. Value Head (Optional): An MLP that takes the current hidden state $h_t$ as input and outputs the value estimate $V(h_t)$. This is common in actor-critic algorithms.

The observational dropout mechanism is applied before the observation encoder. During data collection (rollouts), at each step, a random number is drawn. If it is less than $p$, the input to the observation encoder is replaced with the null observation representation. The agent then proceeds with its forward pass using this potentially nulled input, selects an action, and interacts with the environment. The subsequent training update (e.g., calculating policy gradients or value targets) uses the collected trajectory, including the modified observations $\hat{o}_t$.

def agent_step(h_prev, a_prev, o_current, dropout_prob): """Performs one step of agent interaction with observational dropout.""" # Apply observational dropout if random.random() < dropout_prob: o_hat_current = NULL_OBSERVATION # e.g., zero tensor else: o_hat_current = o_current # Encode observation (or null observation) e_current = observation_encoder(o_hat_current) # Concatenate embedding and previous action rnn_input = concatenate(e_current, action_embedding(a_prev)) # Update recurrent state (implicit world model) h_current = rnn_core(h_prev, rnn_input) # Select action based on hidden state action_distribution = policy_head(h_current) a_current = sample_action(action_distribution) # (Optional) Estimate value value_estimate = value_head(h_current) return a_current, h_current, value_estimate #

The choice of the dropout probability $p$ is a crucial hyperparameter. A value of $p=0$ recovers the standard model-free RL agent. As $p$ increases, the agent is forced to rely more heavily on its internal state. However, excessively high $p$ might hinder learning by starving the agent of essential environmental information. The optimal $p$ is likely environment-dependent.

Emergence of the World Model

The paper posits that the recurrent state $h_t$ implicitly learns to model the environment because doing so is beneficial for maximizing reward under observational dropout. When an observation $o_t$ is dropped ($\hat{o}_t = \emptyset$), the agent must predict relevant aspects of the environment state based on its history, encoded in $h_t$, to select an appropriate action $a_t$. For instance, in a navigation task, $h_t$ might need to encode the agent's estimated position and velocity, or the presence of nearby obstacles, even when direct visual confirmation is temporarily unavailable.

Unlike explicit world models (e.g., MDN-RNN, VAE-based models) that are trained with supervised losses like next-state prediction or reconstruction error, the model emerging from observational dropout is shaped solely by the RL objective. This means the internal representation $h_t$ prioritizes encoding information that is decision-relevant under the observation uncertainty imposed by dropout, rather than accurately predicting all aspects of the next observation. It learns aspects of the world dynamics necessary to bridge the observational gaps and maintain policy performance.

Evidence for the emergence of a functional world model is typically indirect, based on:

1. Improved Performance/Sample Efficiency: Demonstrating that agents trained with observational dropout achieve higher rewards or learn faster than baseline model-free agents (with $p=0$), especially in tasks requiring memory or implicit prediction.
2. Robustness: Showing that the agent can maintain reasonable performance even when observations are temporarily completely withheld during evaluation.
3. Probing/Visualization: Analyzing the internal states $h_t$ to show they correlate with true environment states or dynamics, although this is not the primary focus compared to performance gains.

Experimental Results and Analysis

The experiments were conducted on environments like CarRacing-v0 from OpenAI Gym and tasks within the ViZDoom framework. These environments feature continuous state spaces (pixels) and require memory and understanding of dynamics (e.g., momentum in CarRacing, enemy movement patterns in Doom).

Key findings typically include:

• Performance Improvement: Agents trained with an appropriately chosen $p > 0$ often outperform the baseline model-free agent ($p=0$) in terms of final score and/or sample efficiency on the tested tasks. For example, in CarRacing, which requires anticipating turns and maintaining momentum, observational dropout can lead to significantly better driving policies compared to a purely reactive agent or a standard recurrent agent without dropout.
• Dependence on Dropout Rate: Performance is sensitive to the value of $p$. There usually exists an optimal range for $p$; too low, and the effect is negligible; too high, and learning becomes unstable due to excessive information loss. The paper likely reports results across various $p$ values, showing this dependency. For instance, values of $p$ around 0.1 to 0.3 might be shown to be effective.
• Implicit Prediction: The results suggest that the agent's internal state learns to perform short-term predictions to fill in the gaps caused by dropout. This is inferred from the improved performance on tasks where such implicit prediction is beneficial.
• Comparison to Explicit Models: While not necessarily outperforming state-of-the-art methods based on explicit forward prediction losses (which are directly optimized for prediction accuracy), the observational dropout approach demonstrates that useful world models can emerge without such explicit supervision, offering a simpler alternative mechanism.

Discussion and Practical Implications

The primary implication of this work is that the strong inductive bias provided by explicit predictive modeling might not be strictly necessary for learning useful internal models. By manipulating the agent's information access through observational dropout, the RL objective itself can guide the formation of internal representations that capture environmental dynamics relevant to the task.

Advantages:

• Simplicity: Avoids the need to design and implement potentially complex auxiliary losses for forward prediction or state reconstruction. The modification only involves changing the data stream during rollouts.
• Task-Focused Models: The emergent model is optimized implicitly by the RL objective, potentially leading to representations focused on task-relevant dynamics rather than reconstructing every detail of the observation space.
• Potential for Robustness: Training under observation scarcity might inherently lead to policies that are more robust to noisy or missing sensor data in deployment.

Limitations and Considerations:

• Indirect Control: There is less direct control over what the internal model learns compared to explicit modeling approaches. The learned dynamics are implicit and task-dependent.
• Hyperparameter Sensitivity: The effectiveness relies on tuning the dropout probability $p$, which may vary across environments and tasks.
• Scalability: It remains an open question how well this implicit approach scales to highly complex environments requiring long-term, high-fidelity prediction compared to methods that explicitly optimize predictive accuracy.
• Interpretability: Understanding precisely what the emergent world model represents can be challenging, similar to interpreting hidden states in any RNN.

Applications:

This technique could be valuable in scenarios where:

• Implementing complex predictive models is challenging or computationally expensive.
• The primary goal is robust policy performance under potential sensor intermittency, rather than accurate state prediction itself.
• Simulating biological learning processes where explicit predictive targets are unlikely is of interest.

Future research could explore adaptive dropout schedules, combining observational dropout with other representation learning techniques, or applying it to more complex, long-horizon tasks.

Conclusion

"Learning to Predict Without Looking Ahead" introduces observational dropout as a simple yet effective method for implicitly encouraging reinforcement learning agents to develop internal world models. By stochastically withholding observations, the agent is forced to rely on its internal recurrent state, which consequently learns to capture task-relevant environmental dynamics without being trained on an explicit forward-predictive loss. The experimental results demonstrate the potential of this approach to improve agent performance and sample efficiency in challenging control tasks, offering a compelling alternative perspective on how predictive world models can be acquired.