Policies Learned in Imagination › Worked Example: A One-Dimensional Point Mass

Abstract warnings about exploitation are easy to nod along to and hard to internalize. A concrete example makes the failure visible.

Our environment is a point mass with drag moving on a line. The state is position and velocity, $s = (x, v)$, and the action $a$ is an acceleration. The dynamics are linear:

$$x_{t+1} = x_t + \Delta t\, v_t, \qquad v_{t+1} = \rho\, v_t + \Delta t\, a_t,$$

with $\Delta t = 0.1$ and $\rho = 0.99$. The reward penalizes distance from a target at $x^\star = 2.0$ and the squared control effort:

$$r(s, a) = -(x - x^\star)^2 - 0.05\,a^2.$$

The maximum per-step reward is $0$ when $x=x^\star$ and $a=0$, regardless of velocity. Position, velocity, action, and reward are normalized toy quantities rather than physical SI measurements. Evaluation runs for 40 steps from states near rest at the origin.

The system is linear, low-dimensional, and fully observed. That removes representation learning and partial observability from this particular demonstration. The learned MLP is still an approximation, and optimizer behavior can interact with its errors. We will measure both prediction and closed-loop control rather than attribute every outcome to coverage alone.

The comparison uses two world models trained on different datasets.

• Model A is trained on data collected from initial states near the origin, with position and velocity independently drawn from zero-mean normals with standard deviation $0.25$, and with small actions. It has essentially never seen the region around $x = 2.0$.
• Model B is trained on data collected from initial states spread uniformly over $[-3, 3]$ with wider actions. It has seen the goal region many times.

Both models have the same architecture, transition count, loss, and optimizer. The treatment is the data collection distribution: Model A uses concentrated starts and smaller random actions; Model B uses wider starts and actions. This changes several state-action marginals together, not one isolated scalar.

The comparison can show how those two collection regimes affect learned predictions and downstream policies, but it cannot prove that goal-state coverage alone caused every behavioral difference. Training randomness, action coverage, and the actor optimizer also matter. A small seed check below keeps one attractive trajectory from standing in for a general result.

Model A's reward prediction near and beyond the goal relies much more on extrapolation. The learned policy may therefore see less cost to overshooting than the true reward assigns. Model B has direct samples around the goal, but broader data is not a guarantee of a good controller: the actor still needs a stable objective, bounded actions, and closed-loop evaluation.

In the reproducible run below, the narrow-model policy overshoots the target more than the broad-model policy and has a positive imagination gap. Both overshoot to some extent. That is evidence for this particular setup, not a claim that all narrow models fail or all broad models succeed.

We track the imagination gap: the return the learned model estimates for the policy minus its return in the real environment.