Planning, Control, and Policy Evaluation › Task Return and Sample Efficiency

Task return is the sum of rewards earned by a policy in an episode of finite length $T$, $G(\pi, s_0) = \sum_{t=0}^{T-1} \gamma^{t}\, r(s_t, a_t)$ with $s_0 \sim \mu$, where:

• $G(\pi, s_0)$: the total discounted reward accumulated over the episode starting from initial state $s_0$
• $s_0 \sim \mu$: the initial state is drawn from the stated initial-state distribution $\mu$
• $T$: the finite episode horizon (number of steps)
• $\gamma \in (0, 1]$: the discount factor (here we set $\gamma = 1$, giving an undiscounted finite-horizon objective)
• $r(s_t, a_t)$: the scalar reward received at step $t$ in state $s_t$ under action $a_t$
• $a_t = \pi(s_t)$ for a deterministic policy: randomized MPC additionally depends on its planner stream $\omega$
• $s_{t+1} = f(s_t, a_t)$: the next state produced by the true environment transition dynamics

The discount factor and episode horizon belong to the task definition. Here $\gamma=1$ and $T=30$, so the code sums 30 unweighted rewards with no terminal reward. Changing discount or horizon can change policy rankings, but different reward timing does not force a ranking reversal. Comparisons should use the same objective, or explicitly present the different objectives as different tasks.

For a randomized controller, write the expected evaluation return as $J(\pi)=\mathbb{E}_{s_0\sim\mu,\,\omega}[G(\pi,s_0;\omega)]$, where $\omega$ contains the planner's random draws. A stochastic environment would add transition randomness to that expectation. The deterministic PD controller needs no $\omega$; random-shooting MPC does. Our finite estimates use 20 held-out initial states and one prescribed stream per case, seeded $110+i$ for case $i$. Reconstructing those streams pairs comparisons. These estimates do not average over independent training runs or repeated planner streams at each state, and geometric state coverage alone does not establish estimation precision.

An expected-return claim must state its initial-state distribution $\mu$ and any other randomness being averaged. Other legitimate policy claims concern success probability, tail risk, safety, or latency and require their own defined quantities. Different initial-state distributions answer different return questions unless the comparison explicitly accounts for them.

Initial states influence the trajectories and model queries a policy generates. Starting near the origin does not guarantee that a controller stays there, nor does a farther starting state guarantee worse return. Measure visitation and return if those mechanisms are the claim. Here a fixed uniform sample on $[-1,1]^2$ defines the held-out cases for all source controllers and the paired target experiment.

Sample efficiency describes performance as a function of real-environment interaction budget, not a raw ratio $J/N$. That ratio is especially unhelpful for negative rewards, where its ordering can change simply by changing the reward offset. A learning curve evaluates systems at several disclosed interaction budgets, using either cumulative online learning or separate fits. Independent seeds and uncertainty estimates are needed to characterize variability; their number should be justified by the precision of the comparison. Track these interaction buckets separately:

• Source pretraining interactions: data used before the target adaptation experiment. Source and target can share an environment; the domains must be stated.
• Target adaptation interactions. Real trajectories collected in the deployment environment and used to update any parameter (model, policy, value, planner hyperparameter).
• Validation rollouts. Interactions used to select among hyperparameters or checkpoints. This is not free: those rollouts also consumed actions.
• Test rollouts. Interactions used to report the metric, with no update or selection step consuming their outcomes.

Separating the counters prevents two accounting errors: treating synthetic model transitions as newly collected environment data, and omitting interactions used for selection. Report each bucket so readers can compare the resources their deployment constrains.

Synthetic transitions generated by $\hat f$ are compute, not real samples. A planner that rolls the model $K \cdot H$ times per decision did not collect $K \cdot H$ real environment steps. It spent $K \cdot H$ model evaluations. Both quantities matter, but they belong to different budgets. Real interactions may constrain a deployment when collecting data is expensive or dangerous; model evaluations may constrain it when inference is slow or costly. A method can be sample-efficient and compute-inefficient, or the reverse. Report both budgets against the deployment's actual limits.

A horizon sweep can mix longer lookahead, fewer candidates, and longer model rollouts. A candidate-count sweep changes search adequacy at a fixed model. Neither alone establishes a data-efficiency advantage. To measure that advantage, vary the real-interaction budget and evaluate the resulting learner or checkpoint under a consistent test protocol.

The remaining experiments use a deterministic damped double integrator with two state variables and a known nonlinear transition. This keeps return and resource accounting inspectable. The same evaluation questions apply when studying ensemble methods, imagination-based policies, or value-equivalent models, but those methods need their own implementations and protocols; this notebook does not benchmark them.

The testbed is a deterministic point mass with position and velocity state, damped dynamics, and nonlinear quadratic drag. Its reference transition is a known formula, so we can substitute it into the same planner without claiming to solve control optimally. A neural surrogate would not remove that reference formula; an unknown physical plant would make exact-dynamics substitution unavailable.

The system is a 1D damped point mass with quadratic drag. The state is $s=(p,v)\in\mathbb{R}^2$, the action is a scalar bounded acceleration $a\in[-1,1]$, and its discrete source transition is $v_{t+1}=(1-d)v_t+\Delta t\,g\,a_t-c_{\mathrm{drag}}v_t|v_t|$, followed by $p_{t+1}=p_t+\Delta t\,v_{t+1}$. These are the toy's discrete update rules, not an exact continuous-time solution.

where:

• $v_t$: velocity at time $t$
• $p_t$: position at time $t$
• $a_t$: scalar bounded acceleration applied at time $t$
• $\Delta t$: integration timestep, $0.1$
• $d$: linear damping coefficient, $0.05$
• $c_{\mathrm{drag}}$: quadratic drag coefficient in the discrete velocity update, $0.1$; it is distinct from the return discount $\gamma$
• $g$: actuator gain, set to a source value of $1.0$ and (for transfer) a target value of $0.7$
• $|v_t|$: absolute velocity, capturing the sign-symmetric drag effect

The first line updates velocity from linear damping, the scaled control input, and quadratic drag. The second line updates position by semi-implicit (post-update velocity) Euler integration. The reward is a quadratic tracking-and-control penalty to the origin, $r(s, a) = -\bigl( p^2 + 0.05\, v^2 + 0.01\, a^2 \bigr)$, where:

• $p^2$: position penalty, driving the state toward the origin
• $0.05\, v^2$: velocity penalty, discouraging high speed
• $0.01\, a^2$: control effort penalty, discouraging large actions
• the overall negative sign: penalizes deviations rather than rewarding them, so higher return means staying closer to the origin with less effort

The quadratic term $c_{\mathrm{drag}}v|v|$ makes the true transition nonlinear, while the learned surrogate is affine in state and action. Least squares solves the empirical fitting problem within that class. This is a deliberate class mismatch, not an injected corruption of the fitted coefficients.

Empirical least-squares optimality does not rule out overfitting or finite-sample estimator variation. There is no transition or observation noise here, but the sampled training trajectories still affect the fitted coefficients. The affine class also cannot represent the quadratic drag everywhere. Held-out errors and control tests measure the consequences of that fit; they do not decompose all error into a single bias term.

We collect source training transitions with a PD behavior policy plus action noise, initialized in a narrow region. This creates a restricted collection regime; whether its errors cause exploitation must be tested rather than assumed.

PRETRAIN_INTERACTIONS records the 3000 source training transitions. The later learned-model comparisons share that source fit, and the final accounting separates these records from adaptation, evaluation transitions, and model-search compute.

Fitting a linear model is a two-line least-squares problem. The design matrix concatenates the two state components, the scalar action, and a bias column so the model can absorb any constant offset.

Before trusting this model for planning, we look at its one-step error on held-out transitions drawn from a wider region than training. This is the teacher-forced protocol from Chapter 61: it measures error on these sampled trajectories, not the globally weakest region or every state a planner might visit.

The held-out errors tend to be larger away from the central region. This is consistent with region-dependent model accuracy, but this plot does not measure planner visitation or establish that the planner seeks high-error states. The later decision tests ask whether these prediction errors matter for the policies evaluated.

The near and far averages summarize two groups of held-out transitions. A scatterplot also shows their within-group variation.

The figure below plots prediction-error magnitude against the Euclidean norm of the numerical state coordinates. Position and velocity have different physical units, so this norm is a declared toy-coordinate diagnostic, not a unit-invariant distance. The split at $0.5$ does not define the boundary of training support, and the figure does not measure planner visitation.

The farther group has a larger mean one-step error in this dataset. Similar group means would not prove uniformly accurate predictions, and different means do not establish optimism or exploitation. Error magnitude does not reveal the sign of a return error.

A planner could favor poorly supported trajectories if their predicted rewards are spuriously high, but this plot does not show that mechanism. The decision tests below measure actual return and return-prediction discrepancies; this particular run does not demonstrate optimistic exploitation.

A fixed PD feedback controller is our no-model comparator. It supplies a concrete return benchmark without consuming model-fitting data.

PD does not consult the learned model, so it cannot select actions because that model predicts them favorably. That does not identify why a learned-model controller might lose to PD: exploitation, poor search, horizon choice, or other effects can all matter. Beating this PD controller is the chosen return benchmark for the toy, not a universal condition for a world model to be useful.

We fix a held-out evaluation set of initial states and compute the PD baseline in both the true environment and the learned model. The pair of numbers is the first instance of a pattern we will repeat throughout.

Even for a simple PD controller whose behavior is not chosen by the model, the model's prediction of PD's return differs from reality. In this run the model predicts a lower return, so its error is pessimistic rather than optimistic. Replacing PD with a model-driven planner can change the gap, but neither its sign nor its cause follows from the PD comparison alone.

This is fixed-budget accounting, not a sample-efficiency learning curve. The model fit uses 3000 source transitions. The fixed PD gains require no fitting transitions, but evaluating PD costs 20 episodes of 30 steps, or 600 environment transitions. Compare both performance and disclosed costs; do not call these 600 transitions 600 rollouts or divide the negative return by a sample count.