|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| \documentclass[11pt]{article}
|
|
|
| \usepackage[margin=1in]{geometry}
|
| \usepackage{amsmath,amssymb}
|
| \usepackage{booktabs}
|
| \usepackage{graphicx}
|
| \usepackage{caption}
|
| \usepackage{subcaption}
|
| \usepackage{xcolor}
|
| \usepackage{microtype}
|
| \usepackage[colorlinks=true,linkcolor=black,citecolor=black,urlcolor=blue]{hyperref}
|
|
|
| \graphicspath{{figures/}}
|
|
|
| \newcommand{\Lsup}{\mathcal{L}_{\mathrm{sup}}}
|
| \newcommand{\zg}{z_{G}}
|
|
|
| \title{\bfseries Arrival-and-Hold:\\
|
| Diagnosing and Fixing Horizon-Reset Procrastination\\
|
| in Latent-Space Control with a Frozen World Model}
|
| \author{LeWM $\times$ PushT control experiment}
|
| \date{}
|
|
|
| \begin{document}
|
| \maketitle
|
|
|
|
|
| \begin{abstract}
|
| \noindent
|
| We train a small amortised controller to plan inside the latent space of a
|
| \emph{frozen} JEPA-style world model (LeWM) on the PushT pushing task, and
|
| report a control pathology that inverts a basic expectation from
|
| model-predictive control. Replanning \emph{more often} made the system
|
| \emph{worse}: executing one action block per plan reached the goal on
|
| $50\%$ of held-out episodes, while committing to the entire five-block plan
|
| reached $88\%$ --- a $38$-point gap ($p<10^{-4}$, exact McNemar on paired
|
| episodes). We show the cause is not model error, not optimiser failure, and
|
| not a compute budget: it is the objective. A terminal goal loss
|
| $d(\hat z_H, \zg)$ asks the controller to \emph{arrive exactly at block
|
| $H$}, so every replan resets the deadline and the agent approaches the goal
|
| asymptotically without ever landing. We formalise this as a contraction
|
| recursion $D_{n+1}=cD_n+b$ with a strictly positive fixed point
|
| $D^\ast=b/(1-c)$, and measure $D^\ast$ directly: $0.203$ for a purely
|
| terminal objective, $0.098$ for the original controller. The fix is a
|
| one-line change to the loss --- relabel each training sample with its true
|
| goal offset $q$, penalise the distance \emph{at} $q$, and add a hold term on
|
| every block after $q$. This drives $D^\ast$ down to $0.040$, lifts the
|
| $m{=}1$ success rate from $50\%$ to $94\%$ ($+44$ points, $p<10^{-4}$),
|
| removes the inversion entirely, and reaches parity with a $300\times30$ CEM
|
| planner ($+4$ points, $p=0.69$) while issuing $760\times$ fewer world-model
|
| evaluations per episode than CEM at its best schedule (and $5000\times$
|
| fewer than CEM at the same schedule). We give the full formula inventory
|
| with the purpose and measured effect of every term, an ablation over the
|
| three objective components, and a note on a survivorship confound that makes
|
| the naive cost metric anti-correlated with success ($r=+0.51$).
|
| \end{abstract}
|
|
|
| \tableofcontents
|
| \newpage
|
|
|
|
|
| \section{Setup}
|
|
|
| \subsection{The frozen world model}
|
|
|
| Everything in this report treats the world model as a fixed, non-trainable
|
| oracle. LeWM consists of a ViT-tiny image encoder $E$ (patch size $14$,
|
| $224$px input, $12$ layers, $3$ heads, embedding width $192$) and a
|
| $6$-layer latent \emph{predictor} $P$ with AdaLN action conditioning. Both
|
| are frozen throughout: no gradient ever reaches their parameters. The only
|
| thing we train is a controller that searches in the latent space they
|
| define.
|
|
|
| An observation $o$ becomes a latent $z = E(o) \in \mathbb{R}^{192}$. Given a
|
| latent and an action block $b$, the predictor advances the latent one step:
|
| \begin{equation}
|
| z' = P(z, b).
|
| \label{eq:predictor}
|
| \end{equation}
|
|
|
| \subsection{Action blocking}
|
|
|
| The environment runs at a frameskip of $5$, so a single world-model
|
| transition consumes five raw environment actions. An \emph{action block} is
|
| therefore
|
| \begin{equation}
|
| b \in \mathbb{R}^{10}, \qquad 10 = 5 \text{ raw actions} \times 2 \text{ dims},
|
| \end{equation}
|
| and a plan of horizon $H$ is a stack $b_{1:H} \in \mathbb{R}^{H\times 10}$.
|
| Throughout, $H=5$: one plan covers $25$ raw environment steps.
|
|
|
| \subsection{Data}
|
|
|
| Latents were pre-extracted for $18{,}685$ demonstration episodes
|
| ($2{,}336{,}736$ frames, $192$-dim). Episodes are split with
|
| \texttt{split\_episodes(n, val\_fraction=0.05, seed=0)}, so training and
|
| validation never share an episode. All closed-loop evaluation uses the same
|
| $50$ seeded held-out episodes for every row in every table, which is what
|
| makes the paired statistics in \S\ref{sec:stats} valid.
|
|
|
| \subsection{Notation}
|
|
|
| \begin{center}
|
| \begin{tabular}{ll}
|
| \toprule
|
| symbol & meaning \\
|
| \midrule
|
| $z_t \in \mathbb{R}^{192}$ & latent state at world-model step $t$ \\
|
| $\zg$ & goal latent \\
|
| $b_j \in \mathbb{R}^{10}$ & $j$-th action block of a plan \\
|
| $H = 5$ & plan horizon (blocks) \\
|
| $N = 3$ & context frames given to the controller \\
|
| $K$ & refinement iterations at inference \\
|
| $m$ & blocks executed before replanning \\
|
| $q \in \{1,\dots,H\}$ & true goal offset of a training sample \\
|
| $d_j$ & predicted latent distance at block $j$ \\
|
| $D_n$ & mean goal distance at the $n$-th replan \\
|
| \bottomrule
|
| \end{tabular}
|
| \end{center}
|
|
|
|
|
| \section{The formula inventory}
|
| \label{sec:formulas}
|
|
|
| This section lists every formula used in the experiment, what it is
|
| \emph{for}, and what it measurably \emph{did}. This is the core of the
|
| report: the entire result is a story about which of these terms was wrong.
|
| Sections~\ref{sec:pathology} onward refer back to these equations by number.
|
|
|
|
|
| \subsection{Latent rollout}
|
|
|
| \begin{equation}
|
| \hat z_0 = z_t, \qquad
|
| \hat z_j = P(\hat z_{j-1},\, b_j), \quad j = 1,\dots,H.
|
| \label{eq:rollout}
|
| \end{equation}
|
|
|
| \paragraph{Purpose.} Turn a candidate plan into a predicted latent
|
| trajectory. Because $P$ is frozen and differentiable, the whole rollout is
|
| one differentiable function of $b_{1:H}$, so gradients of any cost defined
|
| on $\hat z_{1:H}$ flow back to the plan --- and, through the controller that
|
| emitted the plan, to the controller weights. This is what makes an amortised
|
| controller possible at all without ever touching the world model.
|
|
|
| \paragraph{Effect.} The rollout is autoregressive, so prediction error
|
| compounds with $j$. This matters later: the cost at block $5$ is a
|
| \emph{less} reliable target than the cost at block $1$, which is one reason a
|
| purely terminal objective (Eq.~\eqref{eq:goalloss} with $\alpha=0$) is
|
| fragile.
|
|
|
|
|
| \subsection{Goal distance}
|
|
|
| \begin{equation}
|
| d_j \;=\; d(\hat z_j, \zg)
|
| \;=\; \frac{1}{D}\bigl\lVert \hat z_j - \zg \bigr\rVert_2^2,
|
| \qquad D = 192.
|
| \label{eq:dist}
|
| \end{equation}
|
|
|
| \paragraph{Purpose.} A scalar "how far from the goal'' signal in latent
|
| space. Dividing by the latent dimension $D$ makes the number comparable
|
| across latent widths and keeps it $O(1)$, which in turn lets a single
|
| $\lambda$ balance it against the support term without retuning.
|
|
|
| \paragraph{Effect.} This is the quantity every objective below is built
|
| from, and the quantity plotted on the $y$-axis of
|
| Figures~\ref{fig:profiles} and~\ref{fig:contraction}. Note it is a
|
| \emph{latent} distance, not task success --- the two are correlated but not
|
| identical, and \S\ref{sec:survivorship} shows a case where they come apart
|
| badly.
|
|
|
|
|
| \subsection{Path weights}
|
|
|
| \begin{equation}
|
| w_j \;=\; \frac{(j/H)^2}{\sum_{i=1}^{H-1} (i/H)^2},
|
| \qquad j = 1,\dots,H-1.
|
| \label{eq:pathw}
|
| \end{equation}
|
|
|
| \paragraph{Purpose.} A normalised weighting over the \emph{intermediate}
|
| blocks of a plan, used by the path term in Eq.~\eqref{eq:goalloss}. The
|
| quadratic ramp deliberately puts almost no weight on early blocks (the agent
|
| should be free to move away from the goal initially if that is what the task
|
| requires) and increasing weight on blocks near the horizon.
|
|
|
| \paragraph{Effect.} Because $w$ grows with $j$, the path term reinforces
|
| rather than counteracts the terminal term's late-arrival preference. This
|
| turns out to be part of the problem, not part of the solution: the path term
|
| softens the pathology but does not remove it (\S\ref{sec:ablation}).
|
| Eq.~\eqref{eq:pathw} is also why setting $\alpha=0$ is so destructive --- it
|
| was supplying the only pressure toward early arrival.
|
|
|
|
|
| \subsection{Goal loss (the original objective)}
|
|
|
| \begin{equation}
|
| \mathcal{L}_{\text{goal}}
|
| \;=\; \underbrace{d_H}_{\text{terminal}}
|
| \;+\; \alpha \underbrace{\sum_{j=1}^{H-1} w_j\, d_j}_{\text{path}},
|
| \qquad \alpha = 0.05.
|
| \label{eq:goalloss}
|
| \end{equation}
|
|
|
| \paragraph{Purpose.} The standard formulation. Reach the goal by the end of
|
| the plan; the small path term is a shaping bonus that discourages wild
|
| excursions on the way.
|
|
|
| \paragraph{Effect --- this is the bug.} Read the terminal term literally: it
|
| says \emph{be at the goal exactly at block $H$}, and says nothing about
|
| blocks $1$ through $H-1$ except through a weight that is largest nearest
|
| $H$. Under a receding horizon, the deadline moves. Every time we replan, $H$
|
| is again five blocks away, so the optimal behaviour under this loss is to be
|
| \emph{five blocks away} from the goal --- forever. The agent procrastinates
|
| by construction. \S\ref{sec:pathology} measures this; setting $\alpha=0$
|
| (pure terminal) makes it dramatically worse, which is the cleanest possible
|
| confirmation that the terminal term is the culprit.
|
|
|
|
|
| \subsection{Arrival-and-hold loss (the fix)}
|
|
|
| Each training sample carries the offset $q$ at which its goal frame actually
|
| occurs. During dataset construction, $q$ is sampled as
|
| \begin{equation}
|
| q \sim \mathcal{U}\{1,\dots,\min(H,\ \text{reach})\},
|
| \end{equation}
|
| where \emph{reach} is how many blocks remain in the episode. The loss is
|
| then
|
| \begin{equation}
|
| \boxed{\;
|
| \mathcal{L}_{\text{ah}}
|
| \;=\; \underbrace{d_q}_{\text{arrival}}
|
| \;+\; \lambda_h \underbrace{\frac{1}{H-q}\sum_{j>q} d_j}_{\text{hold}}
|
| \;}
|
| \label{eq:ahloss}
|
| \end{equation}
|
| (the hold term is defined as $0$ when $q=H$, i.e.\ when there are no blocks
|
| after arrival).
|
|
|
| \paragraph{Purpose of the arrival term.} Penalise the distance at the block
|
| where the goal \emph{actually is}, not at a fixed deadline. This makes the
|
| objective invariant to how far away the goal happens to be, which is exactly
|
| the invariance a receding-horizon controller needs.
|
|
|
| \paragraph{Purpose of the hold term.} Arrival alone is not enough: it says
|
| "be at the goal at block $q$'' but is indifferent to what happens next, so a
|
| controller could sail straight through the goal. The hold term says
|
| \emph{stay there}. It converts the goal from a waypoint into an attractor.
|
|
|
| \paragraph{Critical detail.} $q$ indexes the loss only. It is
|
| \textbf{never} fed to the controller. At inference the controller has no
|
| idea how far the goal is --- it simply learns, over the training
|
| distribution of offsets, to get to the goal as early as possible and stay.
|
| Had we conditioned on $q$, the fix would be a cheat (an oracle input
|
| unavailable at test time) rather than a fix.
|
|
|
| \paragraph{Effect.} The measured consequence is the central result of this
|
| report. In Figure~\ref{fig:profiles}, the terminal objective's distance
|
| profile bottoms out at block $5$ \emph{regardless of $q$}; under
|
| arrival-and-hold, the minimum tracks $q$. Success at $m{=}1$ goes from
|
| $50\%$ to $94\%$, and the contraction fixed point $D^\ast$ falls from
|
| $0.098$ to $0.040$.
|
|
|
| \paragraph{On $\lambda_h$.} We swept $\lambda_h \in \{0, 0.5, 1\}$. All three
|
| remove the pathology; the differences between them are not statistically
|
| distinguishable at $n=50$ (\S\ref{sec:stats}). $\lambda_h=0.5$ is the best
|
| point estimate at $94\%$ and is used as the headline configuration, but the
|
| honest reading is that \emph{the arrival relabelling does the work} and the
|
| hold term is a modest refinement.
|
|
|
|
|
| \subsection{Refinement loss}
|
|
|
| The controller emits a plan and then iteratively refines it $K$ times. All
|
| $K+1$ intermediate plans are supervised, with geometrically increasing
|
| weight:
|
| \begin{equation}
|
| \mathcal{L}_{\text{ref}}
|
| \;=\; \frac{\sum_{k=0}^{K} \rho_k\, \mathcal{L}^{(k)}}
|
| {\sum_{k=0}^{K} \rho_k},
|
| \qquad \rho_k = 2^k,
|
| \label{eq:refloss}
|
| \end{equation}
|
| where $\mathcal{L}^{(k)}$ is Eq.~\eqref{eq:goalloss} or
|
| Eq.~\eqref{eq:ahloss} evaluated on the $k$-th refined plan.
|
|
|
| \paragraph{Purpose.} Two things at once. First, every iterate is a valid
|
| plan, so the controller degrades gracefully if we cut refinement short.
|
| Second, the $2^k$ ramp makes later iterates matter more, which is what
|
| pressures the refinement operator to actually \emph{improve} the plan rather
|
| than just perturb it. Eq.~\eqref{eq:refloss} is the outer wrapper around
|
| whichever inner objective is in use, so swapping
|
| Eq.~\eqref{eq:goalloss} for Eq.~\eqref{eq:ahloss} is genuinely a one-line
|
| change.
|
|
|
| \paragraph{Effect.} Figure~\ref{fig:refinement} shows the cost dropping
|
| sharply over the first three refinements --- and then, past the trained
|
| depth $K=3$, flattening or slightly \emph{rising}. The mean plan change
|
| $\lvert b^{(k)}-b^{(k-1)}\rvert$ decays but never reaches zero, so
|
| refinement is not converging to a fixed point; it is a learned $K$-step
|
| improvement operator, not an optimiser. This is why $K=5$ at inference is
|
| not reliably better than $K=3$ (\S\ref{sec:ablation}).
|
|
|
|
|
| \subsection{Behaviour density and the support term}
|
|
|
| A conditional Gaussian mixture $p_\theta(b \mid c)$ over action blocks
|
| (16 components, width 256, conditioned on the context embedding $c$) is
|
| fit to the demonstration data. Its per-dimension negative log-likelihood is
|
| \begin{equation}
|
| s(c, b) \;=\; -\frac{1}{10}\log p_\theta(b \mid c),
|
| \label{eq:nll}
|
| \end{equation}
|
| and the support penalty is a one-sided hinge against a threshold $c_{95}$:
|
| \begin{equation}
|
| \Lsup \;=\; \mathbb{E}\Bigl[\bigl(\max(0,\; s(c,b) - c_{95})\bigr)^2\Bigr],
|
| \qquad c_{95} = 1.5306.
|
| \label{eq:support}
|
| \end{equation}
|
| $c_{95}$ is the $95$th percentile of $s$ over the demonstration set, so by
|
| construction $5\%$ of real demonstration blocks violate it.
|
|
|
| \paragraph{Purpose.} The world model is only accurate on the action
|
| distribution it was trained on. Without a constraint, a planner optimising
|
| $d_H$ will happily find adversarial action sequences that the predictor
|
| \emph{believes} reach the goal but that the real environment does not
|
| follow. The hinge is one-sided so that being \emph{more} typical than the
|
| threshold is free --- we want to bound exploitation, not clone behaviour.
|
|
|
| \paragraph{Effect.} Removing it ($\lambda_{\text{sup}}=0$) raises the
|
| violation fraction from $0.187$ to $0.652$ --- the controller immediately
|
| drifts off the demonstration manifold. But success is
|
| \emph{unchanged}: $50\%$ vs $50\%$ at $m{=}1$ ($p=1.0$). The term does what
|
| it says, and what it says was not the bottleneck. See
|
| Table~\ref{tab:support} and \S\ref{sec:ablation}.
|
|
|
|
|
| \subsection{Total objective}
|
|
|
| \begin{equation}
|
| \mathcal{L}
|
| \;=\; \mathcal{L}_{\text{ref}}
|
| \;+\; \lambda_{\text{sup}}\, \Lsup,
|
| \qquad \lambda_{\text{sup}} = 0.01.
|
| \label{eq:total}
|
| \end{equation}
|
|
|
|
|
| \subsection{Controller parameterisation}
|
|
|
| The controller conditions on $N{+}1$ tokens (the $N=3$ context latents plus
|
| the goal latent) and emits a plan. Raw plan logits are squashed and rescaled
|
| into the action range:
|
| \begin{equation}
|
| b \;=\; \mu_a \;+\; \sigma_a \odot \tanh(\tilde b),
|
| \label{eq:squash}
|
| \end{equation}
|
| with $\mu_a,\sigma_a$ the per-dimension action mean and standard deviation
|
| of the demonstration set.
|
|
|
| \paragraph{Purpose.} Hard-bound the action range without a clipping
|
| discontinuity, and centre the parameterisation on the data so that
|
| $\tilde b = 0$ is already a reasonable plan.
|
|
|
| \paragraph{Effect.} $\tanh$ saturation means gradients vanish at the
|
| extremes, which is a real cost --- but it makes the plan trivially
|
| environment-safe and removes the need for a separate action-bound penalty.
|
| Eq.~\eqref{eq:squash} also means the support term of
|
| Eq.~\eqref{eq:support} is the only thing constraining \emph{which} in-range
|
| actions the controller may pick.
|
|
|
| Each refinement is a learned residual with a learned, per-iteration step
|
| size:
|
| \begin{equation}
|
| \tilde b^{(k+1)}
|
| \;=\; \tilde b^{(k)} \;+\; \sigma\!\bigl(\gamma_{\min(k, K_{\max})}\bigr)
|
| \cdot \Delta^{(k)},
|
| \label{eq:step}
|
| \end{equation}
|
| where $\sigma$ is the logistic function, $\gamma$ are learned logits, and
|
| $\Delta^{(k)}$ is produced from the features
|
| $[\,\tilde b^{(k)},\; \hat z_{1:H},\; \hat z_{1:H}-\zg,\; d_{1:H}\,]$.
|
|
|
| \paragraph{Purpose.} Giving the refiner the current plan, the predicted
|
| trajectory, the goal residual and the distances is what lets it behave like
|
| a learned gradient step without ever running backpropagation at inference.
|
| The $\sigma(\gamma)$ gate keeps every step in $(0,1)$, so refinement cannot
|
| diverge.
|
|
|
| \paragraph{Effect.} The learned step sizes at $K{=}3$ came out as
|
| $[0.55, 0.45, 0.31]$ --- monotonically decreasing, i.e.\ the controller
|
| learned a decaying schedule on its own. Beyond the trained depth the last
|
| step size is reused, which is exactly why
|
| Figure~\ref{fig:refinement} shows no further improvement past $k=3$.
|
|
|
|
|
| \subsection{Contraction model}
|
| \label{sec:contraction-model}
|
|
|
| To make the pathology quantitative, we model the closed loop as a scalar
|
| affine recursion on the mean goal distance across replans:
|
| \begin{equation}
|
| D_{n+1} \;=\; c\,D_n \;+\; b,
|
| \label{eq:contraction}
|
| \end{equation}
|
| fit by least squares over consecutive replans. If $|c|<1$ this converges to
|
| \begin{equation}
|
| D^\ast \;=\; \frac{b}{1-c}.
|
| \label{eq:fixedpoint}
|
| \end{equation}
|
|
|
| \paragraph{Purpose.} $c$ is the per-replan contraction rate --- how much of
|
| the remaining distance the controller removes per decision. $b$ is the
|
| constant floor it re-introduces each time. The fixed point $D^\ast$ is the
|
| distance at which those two balance: \textbf{the residual error the
|
| closed loop settles at, no matter how long you run it.} A controller with
|
| $b>0$ literally cannot reach the goal.
|
|
|
| \paragraph{Effect.} This is where the diagnosis becomes a number.
|
| Table~\ref{tab:contraction} gives $D^\ast=0.203$ for the terminal-only
|
| objective, $0.098$ for the original, and $0.036$--$0.048$ for the three
|
| arrival-and-hold variants. The ordering matches success rate exactly. Note
|
| also the $R^2$ column: the fit is excellent for terminal-only ($0.938$) and
|
| progressively worse for the corrected controllers ($\approx 0.61$) --- which
|
| is itself informative, because the corrected controllers \emph{terminate}
|
| (they succeed and the episode ends) rather than settling into the smooth
|
| geometric decay that the model describes.
|
|
|
| \paragraph{Caveat.} At $m{=}5$ every fit returns $c>1$ ($1.07$--$1.11$), so
|
| Eq.~\eqref{eq:fixedpoint} yields a negative "fixed point'' and is not
|
| interpretable. We report $m{=}1$ fits only. The $m{=}5$ result is not a
|
| failure of the controller but of the model: with only a handful of replans
|
| per episode, and successful episodes terminating early, the surviving trace
|
| is dominated by the hard episodes and rises.
|
|
|
|
|
| \subsection{Paired statistics}
|
| \label{sec:stats-formulas}
|
|
|
| All rows share the same $50$ seeded held-out episodes, so comparisons are
|
| paired; \S\ref{sec:stats} applies these tests. Let
|
| $a_i, b_i \in \{0,1\}$ be the per-episode outcomes of two planners. Using only the discordant episodes
|
| $n_{01} = \lvert\{i: a_i{=}0, b_i{=}1\}\rvert$ and
|
| $n_{10} = \lvert\{i: a_i{=}1, b_i{=}0\}\rvert$, the two-sided exact McNemar
|
| $p$-value is
|
| \begin{equation}
|
| p \;=\; \min\!\left(1,\;
|
| 2 \cdot 2^{-n} \sum_{i=0}^{k} \binom{n}{i}\right),
|
| \quad n = n_{01}+n_{10}, \quad k = \min(n_{01}, n_{10}).
|
| \label{eq:mcnemar}
|
| \end{equation}
|
| The interval is a percentile bootstrap over episodes ($20{,}000$
|
| resamples, seed $0$) on the paired difference in success rate.
|
|
|
| \paragraph{Purpose.} With $n=50$, one flipped episode moves the success
|
| rate by $2$ points. Treating two rows as independent binomials would ignore
|
| that they are the \emph{same} episodes and badly overstate the uncertainty;
|
| conditioning on the discordant pairs is the correct test.
|
|
|
| \paragraph{Effect.} It changes conclusions. The $+6$-point gap for the
|
| headline controller ($94\%$ vs $88\%$) is \emph{not} significant
|
| ($p=0.25$), and neither is its $+4$-point edge over CEM ($p=0.69$). The
|
| $38$- and $72$-point pathology gaps, by contrast, are overwhelming
|
| ($p<10^{-4}$). Without the paired test one would be tempted to report a
|
| ranking among the three corrected variants that the data does not support.
|
|
|
|
|
| \section{The pathology}
|
| \label{sec:pathology}
|
|
|
| \subsection{The observation}
|
|
|
| The receding-horizon parameter $m$ controls how many of the $H=5$ planned
|
| blocks are executed before replanning. Standard MPC theory says smaller $m$
|
| is better: replanning more often lets the controller correct for model
|
| error, so $m{=}1$ should dominate $m{=}5$.
|
|
|
| It does the opposite (Figure~\ref{fig:sweep}).
|
|
|
| \begin{figure}[htbp]
|
| \centering
|
| \includegraphics[width=0.62\textwidth]{fig1_execution_sweep.pdf}
|
| \caption{Success rate against the number of blocks executed per plan.
|
| Replanning \emph{less} often is monotonically better, for both the learned
|
| controller and a gradient-free CEM planner. The effect is $38$ points for
|
| the controller and $54$ for CEM. That both planners show it rules out an
|
| optimiser bug and points at the shared objective.}
|
| \label{fig:sweep}
|
| \end{figure}
|
|
|
| Crucially, CEM --- which shares the objective but shares no code path with
|
| the controller --- shows the same inversion. That is the observation that
|
| redirected the investigation from the controller to the loss.
|
|
|
| \subsection{The mechanism}
|
|
|
| The terminal loss $d_H$ in Eq.~\eqref{eq:goalloss} asks the controller to be
|
| at the goal \emph{at block $H$}. Under a receding horizon, block $H$ is
|
| always five blocks in the future. The deadline is reset before it is ever
|
| reached, so the controller's learned policy --- approach to a distance that
|
| is optimal to be at \emph{five blocks before arrival} --- is a stable,
|
| self-reinforcing state. It procrastinates.
|
|
|
| Figure~\ref{fig:profiles} is the direct evidence. For each goal offset $q$
|
| in the validation set, we plot the predicted distance
|
| (Eq.~\eqref{eq:dist}) at every block of the rollout
|
| (Eq.~\eqref{eq:rollout}).
|
|
|
| \begin{figure}[htbp]
|
| \centering
|
| \includegraphics[width=\textwidth]{fig2_arrival_profiles.pdf}
|
| \caption{Predicted distance $d_j$ at each plan block, one curve per true
|
| goal offset $q$; stars mark $\arg\min_j d_j$. \textbf{Left} (terminal-only)
|
| and \textbf{middle} (original): the minimum is pinned at block $5$ for
|
| every $q$ --- the controller always plans to arrive at the horizon,
|
| regardless of where the goal actually is. \textbf{Right}
|
| (arrival-and-hold): the minimum tracks $q$, and for $q{=}1$ the profile is
|
| \emph{inverted} --- closest at block $1$, then held. This is the fix
|
| working.}
|
| \label{fig:profiles}
|
| \end{figure}
|
|
|
| \subsection{Quantifying it}
|
|
|
| Fitting Eq.~\eqref{eq:contraction} to the closed-loop traces turns the
|
| qualitative story into a number.
|
|
|
| \begin{figure}[htbp]
|
| \centering
|
| \includegraphics[width=\textwidth]{fig3_contraction.pdf}
|
| \caption{\textbf{Left:} mean latent goal distance against replan index at
|
| $m{=}1$; dotted lines are the fitted fixed points $D^\ast$. The
|
| terminal-only controller plateaus an order of magnitude short of the goal.
|
| \textbf{Right:} the fitted $D^\ast=b/(1-c)$ per variant, annotated with
|
| the underlying $c$ and $b$. The ordering matches success rate exactly.}
|
| \label{fig:contraction}
|
| \end{figure}
|
|
|
| \begin{table}[htbp]
|
| \centering
|
| \caption{Contraction fits at $m{=}1$, Eq.~\eqref{eq:contraction}. $c$ is
|
| the per-replan contraction rate, $b$ the re-introduced floor, and
|
| $D^\ast=b/(1-c)$ the residual distance the closed loop settles at. Lower
|
| $D^\ast$ is better.}
|
| \label{tab:contraction}
|
| \input{tables/contraction}
|
| \end{table}
|
|
|
| The terminal-only controller has $c=0.83$: it removes only $17\%$ of the
|
| remaining distance per replan, and re-adds $b=0.035$ each time. That balance
|
| lands at $D^\ast=0.203$, far outside the success threshold. The corrected
|
| controllers roughly halve $c$ \emph{and} shrink $b$, giving
|
| $D^\ast \approx 0.04$.
|
|
|
| \subsection{It shows up during training}
|
|
|
| The pathology does not require closed-loop rollout to detect. Because the
|
| arrival distance $d_q$ is cheap to log alongside the terminal distance
|
| $d_H$, the divergence is visible in the training curves
|
| (Figure~\ref{fig:training}).
|
|
|
| \begin{figure}[htbp]
|
| \centering
|
| \includegraphics[width=0.66\textwidth]{fig8_training_signal.pdf}
|
| \caption{Running validation distances during training. The terminal-only
|
| run drives $d_H$ to $0.013$ while its arrival cost $d_q$ \emph{rises} to
|
| $0.199$ --- a $13\times$ gap. The corrected run keeps the two within
|
| $1.5\times$ of each other. Monitoring both is a cheap early-warning
|
| signal: a widening gap means the controller is learning to arrive late.
|
| The sawtooth at steps $5000$ and $10000$ is the horizon curriculum
|
| stepping from $2\to3\to5$ blocks.}
|
| \label{fig:training}
|
| \end{figure}
|
|
|
| \begin{table}[htbp]
|
| \centering
|
| \caption{Training configuration and final validation losses. All runs:
|
| $20{,}000$ steps, batch $128$, Adam at $3\times10^{-4}$, weight decay
|
| $10^{-4}$, width $256$, depth $4$, $8$ heads, dropout $0.1$, $K=3$
|
| refinements, horizon curriculum \texttt{0:2, 0.25:3, 0.5:5}. Controller
|
| size: $6.80$M parameters.}
|
| \label{tab:training}
|
| \input{tables/training}
|
| \end{table}
|
|
|
| Note the counterintuitive row ordering in Table~\ref{tab:training}: the
|
| terminal-only run has the \emph{best} terminal validation loss ($0.0130$)
|
| and the \emph{worst} task success ($18\%$). It is not underfit. It is
|
| solving the objective it was given, correctly, and that objective is wrong.
|
|
|
|
|
| \section{Results}
|
|
|
| \begin{table}[htbp]
|
| \centering
|
| \caption{Success rate (\%) on $50$ seeded held-out episodes at both
|
| execution schedules, with the $m{=}1$ minus $m{=}5$ gap and planning cost.
|
| \emph{rows/ep} is world-model predictor rows per episode; \emph{rows/call}
|
| is per solver call, which removes the episode-length confound discussed in
|
| \S\ref{sec:survivorship}. Controllers use $K=3$.}
|
| \label{tab:main}
|
| \input{tables/main_results}
|
| \end{table}
|
|
|
| The headline numbers (Table~\ref{tab:main}): the original controller loses
|
| $38$ points by replanning every block. The corrected controller does not ---
|
| it \emph{gains} $6$ --- and its $m{=}1$ success rate of $94\%$ is the best
|
| result in the entire experiment, above both CEM at its best schedule
|
| ($90\%$) and the original controller at its best schedule ($88\%$).
|
|
|
| \subsection{Cost}
|
|
|
| Figure~\ref{fig:pareto} places every configuration on the cost/accuracy
|
| plane.
|
|
|
| \begin{figure}[htbp]
|
| \centering
|
| \includegraphics[width=0.72\textwidth]{fig6_pareto.pdf}
|
| \caption{Success against planning cost (log scale). The corrected
|
| controllers sit at the top-left: highest success, and roughly $760\times$
|
| fewer world-model evaluations per episode than CEM at its best schedule.
|
| The original controller at $m{=}1$ (the $\times$) is strictly dominated ---
|
| it costs twice as much as the corrected controllers because its episodes
|
| run longer, and succeeds half as often.}
|
| \label{fig:pareto}
|
| \end{figure}
|
|
|
| CEM at $m{=}5$ reaches $90\%$ using $55{,}800$ predictor rows per episode.
|
| The corrected controller reaches $94\%$ using $73$ --- a $760\times$
|
| reduction, and $5000\times$ against CEM at the same $m{=}1$ schedule. The
|
| difference in success is not statistically significant ($+4$ points,
|
| $p=0.69$); the difference in cost is between two and three orders of
|
| magnitude. Wall-clock tells the same story: $0.26$ s per episode against
|
| $1.98$ s. That is the practical case for amortising the planner --- but only
|
| once the objective is right, since the \emph{original} amortised controller
|
| was worse than CEM at $m{=}5$ despite the same cost advantage.
|
|
|
|
|
| \section{Ablation study}
|
| \label{sec:ablation}
|
|
|
| The objective, Eq.~\eqref{eq:total}, has three components beyond the
|
| terminal term: the path term ($\alpha$), the support term
|
| ($\lambda_{\text{sup}}$), and the arrival/hold relabelling ($\lambda_h$). We
|
| ablate each.
|
|
|
| \begin{figure}[htbp]
|
| \centering
|
| \includegraphics[width=0.86\textwidth]{fig4_ablation.pdf}
|
| \caption{Success at both execution schedules for every objective variant.
|
| The number below each pair is the gap ($m{=}1$ minus $m{=}5$): red is the
|
| pathology, green is its absence. Every variant reaches $88$--$92\%$ at
|
| $m{=}5$ --- the differences are entirely in the $m{=}1$ column, which is
|
| precisely the claim that the objective, not the model or the capacity,
|
| determines closed-loop behaviour.}
|
| \label{fig:ablation}
|
| \end{figure}
|
|
|
| \subsection{Path term ($\alpha: 0.05 \to 0$)}
|
|
|
| Removing the path term is the most destructive single change:
|
| $50\% \to 18\%$ at $m{=}1$ ($-32$ points, $p=0.0004$), and the gap widens
|
| from $-38$ to $-72$ (Figure~\ref{fig:ablation}, leftmost pair). The
|
| contraction rate degrades from $c=0.58$ to $c=0.83$ and $D^\ast$ doubles.
|
|
|
| \paragraph{Reading.} The path term was \emph{partially masking} the
|
| pathology. Because $w_j$ weights blocks near the horizon most, it applies
|
| some pressure to be close to the goal before block $H$ --- a weak, indirect
|
| version of the arrival term. Removing it exposes the terminal objective in
|
| its pure form. This is the ablation that identified the terminal term as the
|
| root cause: if the path term helps by pulling the cost earlier, then the
|
| problem is that the cost is too late.
|
|
|
| \subsection{Support term ($\lambda_{\text{sup}}: 0.01 \to 0$)}
|
|
|
| \begin{table}[htbp]
|
| \centering
|
| \caption{Support statistics at $m{=}1$, from Eq.~\eqref{eq:nll} and
|
| Eq.~\eqref{eq:support}. Violation fraction is the share of emitted blocks
|
| with NLL/dim above $c_{95}$. By construction $5\%$ of \emph{demonstration}
|
| blocks exceed the threshold.}
|
| \label{tab:support}
|
| \input{tables/support}
|
| \end{table}
|
|
|
| Removing the support term does exactly what it should to the density
|
| statistics --- the violation fraction jumps from $0.187$ to $0.652$, and
|
| $\Lsup$ rises $12\times$ --- and does \emph{nothing} to task success: $50\%$
|
| vs $50\%$ at $m{=}1$ ($\Delta = 0$, $p = 1.0$), $90\%$ vs $88\%$ at $m{=}5$.
|
|
|
| \paragraph{Reading.} This is a genuine negative result and worth stating
|
| plainly. On this task the world model is evidently robust enough that
|
| off-manifold actions do not produce exploitable prediction error at the
|
| scale the controller can find. We keep the term because it is nearly free
|
| ($\lambda_{\text{sup}} = 0.01$) and because the failure mode it guards
|
| against is catastrophic when it does occur --- but on PushT with LeWM, it is
|
| insurance, not a load-bearing component. Note also that the corrected
|
| controller's violation fraction ($0.208$) is slightly \emph{higher} than the
|
| original's ($0.187$): arriving early requires more decisive action blocks,
|
| which sit further into the tail of the demonstration distribution.
|
|
|
| \subsection{Arrival-and-hold ($\lambda_h$)}
|
|
|
| This is the fix. Relabelling the loss to the true goal offset lifts $m{=}1$
|
| success from $50\%$ to $90$--$94\%$ across all three $\lambda_h$ settings
|
| and eliminates the execution-length inversion in every case.
|
|
|
| \begin{center}
|
| \begin{tabular}{lrrr}
|
| \toprule
|
| & $\lambda_h=0$ & $\lambda_h=0.5$ & $\lambda_h=1$ \\
|
| \midrule
|
| success, $m{=}1$ (\%) & 90 & \textbf{94} & 92 \\
|
| gap ($m{=}1 - m{=}5$) & $-2$ & $+6$ & $+4$ \\
|
| $D^\ast$ & 0.0475 & 0.0397 & 0.0362 \\
|
| val $d_q$ & 0.0125 & 0.0167 & 0.0130 \\
|
| \bottomrule
|
| \end{tabular}
|
| \end{center}
|
|
|
| \paragraph{Reading.} The arrival relabelling carries the effect; the hold
|
| term is a refinement. None of the pairwise differences among the three is
|
| significant (all $p \geq 0.62$, \S\ref{sec:stats}), so we do not claim
|
| $\lambda_h=0.5$ is \emph{the} right value --- only that it is the best point
|
| estimate and that any $\lambda_h \in [0,1]$ works. What \emph{is} significant
|
| is all three against the original ($+40$ to $+44$ points, $p<10^{-4}$).
|
|
|
| \subsection{Refinement depth $K$}
|
|
|
| \begin{figure}[htbp]
|
| \centering
|
| \includegraphics[width=\textwidth]{fig5_refinement.pdf}
|
| \caption{\textbf{Left:} terminal ($d_H$, solid) and arrival ($d_q$, dashed)
|
| cost against refinement index; shading marks depth beyond the trained
|
| $K=3$. Note the terminal-only controller's dashed curve sits $21\times$
|
| above its solid one and barely moves --- refinement optimises the
|
| objective it was given, and that objective ignores $d_q$. The corrected
|
| controller's gap is $5.4\times$. \textbf{Right:} the mean plan change
|
| decays but never reaches zero, so refinement is a learned improvement
|
| operator rather than a converging optimiser.}
|
| \label{fig:refinement}
|
| \end{figure}
|
|
|
| Sweeping $K \in \{0,1,2,3,5\}$ at $m{=}5$ on the original controller gives
|
| $66, 86, 82, 88, 90\%$ --- non-monotonic, and the $K{=}0$ case
|
| (a single feedforward plan, no refinement at all) already reaches $66\%$.
|
| Refinement helps, but it is not where the leverage is: changing the
|
| objective moved success by $44$ points, while adding three refinement
|
| iterations moved it by $22$. The learned step sizes of
|
| Eq.~\eqref{eq:step} came out nearly identical in both the original and
|
| corrected runs ($[0.55, 0.45, 0.31]$ vs $[0.55, 0.45, 0.32]$), which is
|
| further evidence that the refinement machinery was never the problem --- the
|
| two controllers refine in the same way, toward different objectives.
|
|
|
|
|
| \section{Statistical validation}
|
| \label{sec:stats}
|
|
|
| \begin{table}[htbp]
|
| \centering
|
| \caption{Paired comparisons on the same $50$ held-out episodes. $\Delta$
|
| is the difference in success rate (percentage points), the interval is a
|
| $20{,}000$-resample percentile bootstrap, and $p$ is the two-sided exact
|
| McNemar test of Eq.~\eqref{eq:mcnemar}. $\ast$ marks $p<0.05$.}
|
| \label{tab:paired}
|
| \input{tables/paired_stats}
|
| \end{table}
|
|
|
| Three things are worth drawing out of Table~\ref{tab:paired}.
|
|
|
| \paragraph{The pathology is real and large.} Every $m{=}1$ vs $m{=}5$
|
| comparison for an uncorrected objective is significant at $p<10^{-4}$, with
|
| confidence intervals that exclude zero by a wide margin. This is not a noise
|
| artefact of $n=50$.
|
|
|
| \paragraph{The fix is real and large.} Arrival-and-hold beats the original
|
| by $+44$ points at $m{=}1$ ($[+30,+58]$, $p<10^{-4}$) and beats
|
| terminal-only by $+76$.
|
|
|
| \paragraph{The fine-grained rankings are not.} The corrected controller's
|
| $+6$ over its own $m{=}5$ schedule ($p=0.25$), its $+4$ over CEM
|
| ($p=0.69$), and all three pairwise $\lambda_h$ comparisons ($p \geq 0.62$)
|
| are indistinguishable from noise. The one marginal result is
|
| arrival-and-hold versus the original at $K{=}3$, $m{=}4$: $+14$ points,
|
| $p=0.039$ --- significant, but only just, and it would not survive a
|
| multiple-comparison correction across the $17$ tests in this table. We
|
| report it as suggestive rather than established.
|
|
|
|
|
| \section{A survivorship confound}
|
| \label{sec:survivorship}
|
|
|
| One measurement in this experiment is actively misleading, and it is worth
|
| recording because it nearly inverted a conclusion.
|
|
|
| PushT episodes \textbf{terminate on success}
|
| ($\lVert\Delta \text{pos}\rVert < 20$ and
|
| $\lvert\Delta\theta\rvert < \pi/9$). The number of solver calls per
|
| evaluation is fixed by the schedule ($200$ at $m{=}1$, $40$ at $m{=}5$), but
|
| the number of environments \emph{still running} at each call is not: good
|
| controllers finish early and drop out.
|
|
|
| Consequently:
|
| \begin{equation}
|
| \frac{\text{predictor rows}}{\text{solver call}}
|
| \;=\; \text{mean number of episodes still alive},
|
| \label{eq:survivorship}
|
| \end{equation}
|
| which measures \emph{episode length}, not per-decision cost. Every
|
| controller in this study has \emph{identical} per-decision cost --- same
|
| architecture, same $K$, same horizon. Figure~\ref{fig:survivorship}
|
| decomposes this.
|
|
|
| \begin{figure}[htbp]
|
| \centering
|
| \includegraphics[width=\textwidth]{fig7_survivorship.pdf}
|
| \caption{\textbf{Left:} across all $36$ paired evaluation rows, better
|
| controllers report \emph{higher} mean terminal distance ($r=+0.51$).
|
| \textbf{Right:} rows-per-call decomposed --- it is exactly the mean number
|
| of surviving episodes. The terminal-only controller looks $2.5\times$ more
|
| "expensive'' than the corrected one purely because its episodes never
|
| end.}
|
| \label{fig:survivorship}
|
| \end{figure}
|
|
|
| The left panel is the sharper warning. \emph{Mean terminal distance is
|
| positively correlated with success rate} ($r=+0.51$): the better the
|
| controller, the worse its average reported cost. The reason is the same ---
|
| successful episodes exit the average early, leaving the mean dominated by
|
| the hard episodes that a good controller is still working on, while a bad
|
| controller's easy-but-unfinished episodes keep its average low.
|
|
|
| \paragraph{Practical rule.} On any benchmark with success-triggered
|
| termination, per-step cost and per-step error averages are survivorship
|
| statistics. Report cost per \emph{decision} and success separately, and
|
| never rank controllers by mean episode cost.
|
|
|
|
|
| \section{Limitations}
|
|
|
| \begin{itemize}
|
| \item \textbf{$n=50$.} The held-out set is small. It is large enough to
|
| establish the $38$-, $44$- and $72$-point effects with certainty, and far
|
| too small to rank the three corrected variants against each other. We have
|
| been explicit about which claims fall on which side of that line.
|
|
|
| \item \textbf{One task, one world model.} PushT with LeWM. The
|
| horizon-reset argument is a property of the \emph{objective} under a
|
| receding horizon and should generalise, but that is an argument, not
|
| evidence.
|
|
|
| \item \textbf{Single seed per configuration.} Each row is one training run
|
| evaluated on $50$ episodes. The pairing controls episode-level variance,
|
| not seed-level variance.
|
|
|
| \item \textbf{The $m{=}5$ contraction fits are uninterpretable}
|
| ($c>1$), as noted in \S\ref{sec:contraction-model}. The contraction
|
| analysis is evidence at $m{=}1$ only.
|
|
|
| \item \textbf{The support term is unvalidated on this task.} It has the
|
| intended effect on the density statistics and no measurable effect on
|
| success. We cannot say from this experiment whether it would matter on a
|
| task where the world model is more exploitable.
|
|
|
| \item \textbf{Marginal results flagged.} The $p=0.039$ comparison against
|
| $m{=}4$ would not survive correction for the $17$ tests reported.
|
| \end{itemize}
|
|
|
|
|
| \section{Conclusion}
|
|
|
| A learned latent-space controller failed in a way that looked like a
|
| capacity or optimiser problem and was neither. It was a specification
|
| problem: $d(\hat z_H, \zg)$ means \emph{arrive at block $H$}, and under a
|
| receding horizon block $H$ never arrives. The controller learned the correct
|
| solution to the wrong question, which is why its \emph{training} loss was
|
| excellent ($0.0130$, the best of any variant) while its success rate was the
|
| worst ($18\%$).
|
|
|
| The fix required no architectural change, no additional compute, and no
|
| world-model retraining --- only relabelling the loss to the goal's true
|
| offset $q$ and adding a hold term after it, Eq.~\eqref{eq:ahloss}. The
|
| result is $94\%$ at the most frequent replanning schedule, parity with a CEM
|
| planner using two to three orders of magnitude more world-model calls, and a
|
| contraction fixed point reduced from $0.098$ to $0.040$.
|
|
|
| Three transferable lessons:
|
| \begin{enumerate}
|
| \item \textbf{Under a receding horizon, penalise arrival, not the
|
| terminal step.} A fixed-deadline cost composed with a moving deadline is a
|
| procrastination incentive.
|
| \item \textbf{Log the arrival cost next to the terminal cost.} The gap
|
| between them (Figure~\ref{fig:training}) diagnoses this failure at
|
| training time, with no rollout.
|
| \item \textbf{Beware survivorship in success-terminated benchmarks.} Mean
|
| episode cost and mean episode error both invert.
|
| \end{enumerate}
|
|
|
|
|
| \appendix
|
| \section{Reproducing the report}
|
|
|
| All figures and tables are generated from the raw result files, so nothing
|
| in this document is hand-transcribed:
|
|
|
| \begin{verbatim}
|
| python report/make_figures.py # -> report/figures/*.pdf, *.png
|
| python report/make_tables.py # -> report/tables/*.tex
|
| latexmk -pdf report/report.tex # -> report/report.pdf
|
| \end{verbatim}
|
|
|
| \texttt{make\_figures.py} reads \texttt{data/runs/eval/results.jsonl},
|
| \texttt{data/runs/diagnostics/*} and the saved checkpoints;
|
| \texttt{make\_tables.py} recomputes the paired statistics with the same
|
| functions used during the study (\texttt{scripts/paired\_stats.py}) and
|
| persists them to \texttt{data/runs/eval/paired\_stats.jsonl}.
|
|
|
| The original controller predates the per-offset profile logging, so its
|
| entry in Figure~\ref{fig:profiles} and its arrival column in
|
| Table~\ref{tab:training} come from \texttt{report/recover\_profiles.py},
|
| which recomputes them with the same \texttt{evaluate()} on the same held-out
|
| split and a seeded loader. Its recovered $q{=}1$ profile
|
| $[0.0914, 0.0537, 0.0337, 0.0229, 0.0133]$ matches the values recorded
|
| during the original run, confirming the recovery is faithful.
|
|
|
| \end{document}
|
|
|
