TD-MPC and Control-Centric Representations › Joint Latent Dynamics and Value Learning

A TD-MPC agent has five learned components that share a latent representation. The encoder, dynamics, reward, and value components are optimized with a weighted multistep objective; the policy prior has a separate objective and optimizer. The representation is shaped by control-related losses rather than a pixel-reconstruction loss.

To make this precise, let $o_t$ denote the observation (the raw state vector in a proprioceptive task, an image in a visual one), $a_t \in \mathbb{R}^{m}$ the action, and $r_t$ the scalar reward. The agent maintains a latent state $z_t \in \mathbb{R}^{d}$ and five functions of it:

• Encoder $h_\theta$: maps an observation to a latent state, $z_t = h_\theta(o_t)$.
• Latent dynamics $d_\theta$: predicts the next latent from the current latent and action, $\hat{z}_{t+1} = d_\theta(z_t, a_t)$.
• Reward model $R_\theta$: predicts the immediate reward, $\hat{r}_t = R_\theta(z_t, a_t)$.
• Value function $Q_\theta$: estimates the discounted return from a latent state and an action.
• Policy prior $\pi_\theta$: maps a latent state to an action and supplies noisy proposal trajectories to the planner.

Notice that the five functions divide naturally into two groups. The encoder, dynamics, and reward model describe the world: they say where you are, where you go next, and what you earn along the way. The value function and policy prior make decisions: one estimates how good a state-action pair is, and the other proposes actions worth trying. Joint training is what binds the two groups, because the value loss reaches back through the encoder and forces the world-modeling half to preserve whatever information the decision-making half needs.

The original method has deterministic latent dynamics rather than a belief distribution or stochastic transition model. Sampling candidate actions during planning does not make those latent transitions probabilistic. This distinction matters when the environment has stochastic outcomes.

Training samples sequences of observed transitions from replay, encodes the first observation, and unrolls the latent dynamics along recorded actions. Write $\hat z_0=h_\theta(o_t)$ and $\hat z_{k+1}=d_\theta(\hat z_k,a_{t+k})$. For an unroll of $K$ transitions, the original paper's model objective has three coefficient-weighted losses at each step, discounted in the loss by $\lambda^k$:

$$\mathcal{J}_{\mathrm{model}}= \mathbb{E}_{\mathcal D}\sum_{k=0}^{K-1}\lambda^k\left[ c_r\big(R_\theta(\hat z_k,a_{t+k})-r_{t+k}\big)^2 +c_q\big(Q_\theta(\hat z_k,a_{t+k})-y_{t+k}\big)^2 +c_z\left\|\hat z_{k+1}-\operatorname{sg}\!\left[h_{\bar\theta}(o_{t+k+1})\right]\right\|_2^2 \right].$$

Here $c_r,c_q,c_z$ balance reward prediction, TD value fitting, and latent consistency; they are not equal by definition. $\bar\theta$ denotes slow-moving target parameters. The consistency target encodes the observed next state, while its prediction is reached by rolling forward from the first encoded state. A collapsed constant latent can make consistency easy, but reward and value prediction generally require distinctions between states. Their losses supply task-relevant pressure; none individually guarantees a useful representation.

For example, a constant encoder and matching dynamics achieve zero consistency error, yet cannot distinguish two observations whose rewards or returns differ under the same action. In tasks where those targets vary, minimizing the other losses disfavors that collapse. This is a task-dependent argument, not a proof of convergence or a guarantee against every degenerate local minimum.

The released TD-MPC implementation computes the one-step TD target using the observed successor encoded by the online encoder. With its twin target critics and a noisy action sampled from the current policy, the target is

$$z_{t+1}^{\mathrm{obs}}=h_\theta(o_{t+1}),\qquad \tilde a_{t+1}\sim\pi_\theta(\cdot\mid z_{t+1}^{\mathrm{obs}}),\qquad y_t=r_t+\gamma\min_{j\in\{1,2\}}Q_{\bar\theta,j}(z_{t+1}^{\mathrm{obs}},\tilde a_{t+1}).$$

The current policy samples around its mean with truncated noise for this target; the paper's single-$Q$ notation abstracts the released twin critics. The value loss evaluates both $Q_{\theta,j}(\hat z_k,a_{t+k})$ on predicted latent states during the multistep unroll. The bootstrap target above instead uses the next real observation. The original paper writes its TD target with $z_{k+1}$ in the recurrent loss; reading that notation as a claim that the released code bootstraps from the predicted successor confuses the two paths. A target network slows target drift; it does not guarantee convergence of nonlinear off-policy learning.

The bootstrap action comes from the learned policy rather than an explicit maximization over a continuous action space. This is a tractable actor-critic approximation, and its errors can affect the critic. It does not imply a strict upper bound on critic quality or that an action omitted by the prior is never represented in replay.

The released implementation optimizes the policy separately on the initial and $K$ predicted recurrent latents. It detaches these $K+1$ latents, freezes both critics' parameters for the actor step, and maximizes the smaller critic estimate:

$$\mathcal{J}_{\pi}=-\mathbb{E}_{\mathcal D}\sum_{k=0}^{K}\lambda^k \min_{j\in\{1,2\}}Q_{\theta,j}\big(\operatorname{sg}[\hat z_k],\pi_\theta(\operatorname{sg}[\hat z_k])\big).$$

This writes the released update at the level of its mean objective and omits the action sampling noise. The detached latents prevent this actor step from updating the encoder or dynamics, while the frozen critics still pass gradients through the policy's action. Neither an entropy bonus nor an action-magnitude penalty appears in the original objective. The small code example below adds an explicit action-magnitude penalty as a toy stabilizer and labels it accordingly.

The toy code below includes an episodic termination mask, $(1-\mathrm{terminated}_t)$, but not a time-limit truncation mask. That is a sensible convention for its finite episodes, not a term shown in the original paper's displayed TD equation. Whether a timeout should bootstrap depends on the environment's observation and episode semantics; confusing timeout with a true terminal state can bias fitted values.

Dropping the decoder is not an oversight. Three arguments motivate it, and they are worth stating because they apply well beyond this one algorithm.

Sufficiency beats fidelity. If the robot's arm color carries no information about the task, dynamics, or reward, a control policy need not preserve it. It does need quantities that change the return: joint angles, velocities, contacts, and distances. Color could matter, however, if it identifies a task-relevant object or instruction. Reconstruction asks the latent to retain visual detail even when that detail is irrelevant to control. A decoder-free objective does not impose that demand. This is the same argument that motivates the predictive-embedding family covered in Predictive Embeddings, JEPA, and Energy-Based Models.

The gradient budget goes to task-related predictions. The encoder is trained through reward, latent consistency, and value losses rather than a pixel decoder. This avoids allocating a direct reconstruction loss to visual detail irrelevant to the reward. It does not prove that every retained feature matters for control, nor that reconstruction would dominate every alternative model.

It changes what you monitor. Low reconstruction error alone cannot certify good decisions. Removing the decoder also removes that potentially misleading metric, but a falling value loss is no guarantee of good control either: the critic can fit biased targets or be exploited by the planner. Closed-loop evaluation remains necessary.

Reconstruction can nevertheless provide a useful auxiliary signal when reward and value targets are weak or sparse. The original paper evaluates decoder-free learning on continuous-control tasks, including sparse-reward cases; it does not establish a universal dense-versus-sparse boundary. Whether observation reconstruction helps is an empirical question for the task and data regime.

The original TD-MPC model and TD-MPC2 should not be conflated here. The following normalization choices are associated with TD-MPC2, not the original paper's encoder and dynamics.

Layer normalization is used throughout TD-MPC2's MLP components. The released original TD-MPC implementation does not place it in the dynamics and reward networks as previously claimed. Normalization can improve optimization stability, but its exact effect cannot be inferred from one layer in isolation.

Simplicial normalization (SimNorm), introduced in TD-MPC2, splits the latent vector into fixed-size groups and applies a softmax within each group. Each coordinate lies in $[0,1]$ and each group sums to one. This bounds the coordinates of a normalized latent; it does not guarantee accurate or stable multistep predictions. Our toy below borrows SimNorm from TD-MPC2, so it is an illustrative hybrid rather than a faithful reproduction of the 2022 architecture.

The coordinate bound can be checked on a single deterministic vector. This demonstrates the operation of SimNorm, not a claim about rollout errors.