Planning with Learned Models: Trajectory Sampling

Building an internal model $\hat{M}$ of the environment, typically consisting of a learned transition function $\hat{P}(s'|s, a)$ and a learned reward function $\hat{R}(s, a, s')$, allows agents to gain a predictive understanding of their surroundings. Having such a model, even an approximate one, opens the door to planning. Instead of solely relying on real interactions, the agent can now use its internal model to simulate potential future scenarios and learn from them. One straightforward and versatile way to leverage the learned model is through trajectory sampling, sometimes called simulation-based planning or model-based rollouts.

The core idea is simple: treat the learned model $\hat{M}$ as if it were the real environment and generate sequences of states, actions, and rewards by interacting with it. These simulated trajectories provide "synthetic" experience that can be used to update value functions or policies, often significantly improving sample efficiency compared to purely model-free methods.

Generating Simulated Trajectories

Starting from a given state $s$ (which could be the agent's current state, a state encountered previously, or even a state), the planning process via trajectory sampling unfolds as follows:

1. Select Action: Choose an action $a$ based on some policy. This could be the agent's current learned policy $\pi$, or perhaps a different policy used specifically for planning (e.g., an exploratory policy within the simulation).
2. Predict Outcome: Use the learned model $\hat{M}$ to predict the next state $s'$ and the reward $r$.
   • Sample the next state: $s' \sim \hat{P}(\cdot | s, a)$.
   • Determine the reward: $r = \hat{R}(s, a, s')$ (or sample if the reward function is stochastic: $r \sim \hat{R}(\cdot | s, a, s')$).
3. Repeat: Continue this process from the predicted state $s'$ for a defined number of steps (the planning horizon $H$) or until a termination condition is met within the simulation.

This generates a simulated trajectory: $(s_0, a_0, r_0, s_1, a_1, r_1, \dots, s_H, a_H, r_H)$. Note that $s_0$ is the starting state for the simulation.

A flowchart illustrating the trajectory sampling process using a learned environment model $\hat{M}$. Starting from a state $s$, an action $a$ is chosen, and the model predicts the next state $s'$ and reward $r$. This loop continues for a planning horizon $H$, generating a simulated trajectory used for learning updates.

Using Simulated Trajectories for Planning

Once we have one or more simulated trajectories starting from a state $s$, how can we use them for planning? There are several common approaches:

1. Monte Carlo Evaluation (Value Estimation): Generate multiple trajectories starting from state $s$. The return $G = \sum_{t=0}^{H} \gamma^t r_t$ is calculated for each trajectory. The average of these returns provides an estimate of the state value $V(s)$ under the policy used during simulation. Similarly, if we simulate trajectories starting from state $s$ after taking a specific action $a$, the average return estimates $Q(s, a)$. These estimates can then be used to update the agent's value function approximator (e.g., a neural network). This is precisely the mechanism used in the planning phase of Dyna-Q, where simulated experiences $(s, a, r, s')$ generated by the model are fed into standard Q-learning or SARSA updates.

2. Policy Improvement via Action Selection: For a given state $s$, simulate trajectories for each possible action $a \in \mathcal{A}$. Estimate the Q-value $Q(s, a)$ for each action using the method above. The action with the highest estimated Q-value is then chosen for execution in the real environment. This approach uses the model to "look ahead" and evaluate the consequences of potential actions before committing to one. This resembles Monte Carlo Control applied to the learned model.

3. Generating Data for Policy Optimization: The simulated trajectories $(s_t, a_t, r_t, s_{t+1})$ can be collected into a buffer, similar to an experience replay buffer but populated with synthetic data. This buffer can then be used to train a policy network directly using policy gradient methods or actor-critic algorithms. This allows the policy to be improved based on imagined outcomes.

Approaches to Trajectory Sampling

While powerful, planning with trajectory sampling involves several design choices and potential issues:

• Planning Horizon ($H$): A short horizon $H$ means planning is faster and less prone to accumulating errors from an imperfect model. However, it limits the "depth" of the lookahead. A longer horizon allows for examining more distant consequences but increases computational cost and sensitivity to model inaccuracies.
• Model Accuracy: The effectiveness of planning heavily depends on the quality of the learned model $\hat{M}$. If the model is inaccurate, the simulated trajectories might diverge significantly from reality, leading to suboptimal or even harmful planning outcomes (often referred to as model bias). Techniques exist to quantify model uncertainty and incorporate it into planning, but the fundamental challenge remains.
• Computational Cost: Generating many or long trajectories can be computationally expensive, especially if the model itself is complex (e.g., a large neural network) or the state/action spaces extensive. This often necessitates trade-offs between planning depth/breadth and computational feasibility.
• Choice of Starting States: Where should simulations begin? Common choices include the agent's current real state, states randomly sampled from recent experience, or states identified as particularly important (e.g., states with high prediction error from the model).

"Trajectory sampling provides a flexible bridge between model learning and control. By simulating interactions with its learned world model, the agent can augment its experience, potentially leading to faster learning and better decision-making compared to purely model-free approaches, especially when interactions are costly or limited. The next sections will explore specific architectures like Dyna-Q and planning methods like MCTS that build upon these core ideas."