PAPER REPORTENAll readings ↗

TD-MPC2: Scalable, Robust World Models for Continuous Control

English reading report: Method, equations, original figures, experiments and reproducibility.

Authors: Nicklas Hansen; Hao Su; Xiaolong Wang

Affiliations: University of California San Diego

Source: ICLR 2024 · ref-9fc7047b7d787fa298bf ↗ · Project page ↗ · Catalog record

Reading: 446 / 558 · 6 original figures & tables · ~20 min ·

1. Paper overview

In one sentence: TD-MPC2 makes latent-space planning more robust across continuous-control tasks, while retaining an explicit planner and substantial dependence on task rewards and dataset coverage. e02e03e07e09e11e12e13e18e24

At a glanceWhat to know
Research problem
Source description

Can model-based RL learn from mixed-quality experience across different embodiments, reward scales, and action spaces without fragile per-task tuning? The authors target a limitation of generalist behavior cloning: dependence on near-expert demonstrations. TD-MPC2 instead optimizes expected return using an implicit control model. e02e03

Core mechanism
Source description

A robustness recipe combines simplicial latent normalization, normalized MLPs, discrete reward/value regression, an ensemble of value functions, and a stochastic policy prior. e04e05e06e16

A key reported resultOffline multitask control across 80 tasks: 1M: 16.0 / 3.7; 5M: 49.5 / 4.2; 19M: 57.1 / 5.3; 48M: 68.0 / 12; 317M: 70.6 / 33.

Mean normalized score on a 0–100 scale; approximate GPU days. Fixed 545M-transition dataset; evaluation on the training task set, with increasing model configurations.

Each cost is reported for one NVIDIA GeForce RTX 3090; Figure 7 shows TD-MPC deteriorating as its size increases. Larger TD-MPC2 configurations improve the aggregate score. This is neither unseen-task evaluation nor a fitted scaling law; uncertainty is not provided in Table 1. e11e23

Reading caution
Source description

The experiments are simulated continuous control. Discrete-action planning and broader zero-shot generalization remain open. The authors flag reward misspecification, unconstrained physical autonomy, and data cost as risks. e09e15

Core contributions

  • Source description

    A robustness recipe combines simplicial latent normalization, normalized MLPs, discrete reward/value regression, an ensemble of value functions, and a stochastic policy prior. e04e05e06e16

  • Source description

    Learned task embeddings, zero-padding, and action masking allow shared components to serve tasks with different observation and action dimensions. e08

  • Reader analysis

    The experiments connect single-task learning, offline multitask scaling, and online adaptation, rather than demonstrating unrestricted generalist behavior. e09e11e12

Figure 3. Actions drive latent transitions; reward and value predictions make those transitions useful for control. Original paper, p. 3 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Read each vertical column from observation s through the encoder to latent state z, then follow the horizontal arrows labeled a between latent states. The upward branches denote predicted action, reward, and terminal value. These branches have different roles: the policy prior proposes actions, whereas reward and Q evaluate candidate actions. Equation (2) explicitly conditions dynamics, reward, and Q on an action even where the drawing suppresses those inputs. The later observation encoders supply training targets; they do not provide future observations during planning. Equation (3), rather than an unmarked arrow in this schematic, establishes the stopped gradient through each next-observation target. e03e04e05e07e08

What it supports. The mechanism is a decoder-free control model: it learns a latent rollout that can support return estimation without synthesizing future images. The drawn action prediction is a policy prior, not the complete executed controller. Section 3.2 adds the search process that selects the action used in the environment.

Where the evidence stops. The drawing omits task conditioning, loss wiring, and the planner. Equation (4) also subtracts entropy despite describing maximum-entropy learning; its sign convention remains unresolved. This schematic cannot settle that implementation issue.

2. Motivation

2.1 The problem and the proposed response

Source description

Can model-based RL learn from mixed-quality experience across different embodiments, reward scales, and action spaces without fragile per-task tuning? The authors target a limitation of generalist behavior cloning: dependence on near-expert demonstrations. TD-MPC2 instead optimizes expected return using an implicit control model. e02e03

2.2 What this reading follows

A useful control model need not reconstruct everything the agent sees. TD-MPC2 compresses observations into latent states, predicts the consequences of candidate actions, and uses those predictions to search for a high-return action sequence. The paper's central question is whether this recipe can remain stable as tasks, embodiments, and model size vary. Read the architecture first, then separate three experiments: independently trained online agents, an offline model evaluated across its training tasks, and online adaptation to held-out tasks. The ablations reveal why planning and normalization matter; the appendix shows that transfer gains and offline regularization remain conditional. e02e03e07e09e11e12e13e18e24

3. Research context

We place the paper in the collection through its world–action interface. The catalog labels and the reading’s assessment are shown separately.

Catalog dimensionRecorded classification
Major categoryFoundational work
ArchitectureNot applicable
Prediction paradigmNot applicable
QuadrantNot applicable

3.1 Evidence-based assessment

Supports the recorded classification

Reader analysis

The recorded foundational world-model and planning classification fits the architecture. This is action-conditioned forward latent dynamics with a distinct policy prior, value ensemble, and explicit optimizer. Jointly learning those components does not establish a single unified future/action predictor or inverse-dynamics mechanism, so the recorded not-applicable quadrant remains appropriate. e03e04e07e08

This is the collection’s architectural analysis, not a new related-work survey. Benchmark comparisons and their protocols appear in Section 6.

4. Problem formulation

4.1 Inputs and outputs

InputsOutputs
  • Current proprioceptive state; image observations in the separate visual-RL experiments.
  • Candidate continuous actions; task identity through a learned embedding in multitask models.
  • Replay trajectories containing observations, actions, rewards, and next observations.
  • Predicted next latent state, reward distribution, terminal action value, and policy-prior action.
  • A first action selected by model-predictive control and executed in the environment.

4.2 Equations and their role

L(θ)=E(s,a,r,s)0:HB ⁣[t=0Hλt(ztsg(h(st))22+CE(r^t,rt)+CE(q^t,qt))]\mathcal{L}(\theta)=\mathbb{E}_{(s,a,r,s')_{0:H}\sim\mathcal{B}}\!\left[\sum_{t=0}^{H}\lambda^t\left(\left\|z'_t-\operatorname{sg}(h(s'_t))\right\|_2^2+\operatorname{CE}(\hat r_t,r_t)+\operatorname{CE}(\hat q_t,q_t)\right)\right]
Eq. (3): θ denotes model parameters, B replay, H the rollout horizon, and λ temporal weighting. The predicted next latent z′ matches the stopped-gradient encoding h(s′); CE trains reward and value predictions against soft targets. The implemented loss coefficients are separately specified in Table 8. e04e22
qt=rt+γQˉ ⁣(zt,p(zt))q_t=r_t+\gamma\bar Q\!\left(z'_t,p(z'_t)\right)
The source TD target combines observed reward r with discounted next-state value. γ is the discount factor, p the policy prior, and Q̄ an EMA target value function; the ensemble uses the minimum of two sampled target functions. e04
γ=clip ⁣(T/51T/5,[0.95,0.995]),S=max(5T,1000)\gamma=\operatorname{clip}\!\left(\frac{T/5-1}{T/5},[0.95,0.995]\right),\qquad S=\max(5T,1000)
Appendix H defines T as expected episode length after action repeat and S as seed environment steps before gradient updates. These heuristics qualify the claim of shared hyperparameters; task-specific discounts also apply during multitask learning. e23

5. Method in detail

5.1 Learn the quantities that planning can use

Source description

TD-MPC2 learns from a replay trajectory rather than from a future-image reconstruction target. Starting with an encoded observation, dynamics advances the latent state using the recorded action. The resulting prediction is compared with the stopped-gradient encoding of the actual next observation. Reward and value heads supply complementary supervision: immediate task feedback and a bootstrapped estimate of what follows. The policy prior has its own objective and receives the policy-objective gradients, while the encoder, dynamics, reward, and value functions share the model objective. SimNorm constrains groups of latent coordinates through softmax, and discrete reward/value regression reduces dependence of loss magnitude on reward scale. The source's temperature-limit prose conflicts with its softmax formulation; Table 8's actual default is temperature one, without a temperature sweep. e03e04e05e06e16e22

5.2 Turn prediction into a feedback controller

Source description

At a decision step, the agent knows the current observation but has no future observations available. It therefore rolls candidate actions forward through learned latent dynamics, predicts rewards, and appends a terminal value to account for consequences beyond the short search horizon. MPPI repeatedly samples action sequences and updates the distribution toward promising candidates. The policy prior provides some proposals, while the previous solution helps initialize the next search. Only the first selected action is executed before fresh feedback triggers replanning. This separation explains why the policy branch in Figure 3 cannot be read as the whole controller. The actor ablation compares that branch with planning and their combination. Its reported advantage for planning supports inference-time model use rather than merely an auxiliary prediction objective during training. e03e07e13e22

5.3 Separate task sharing from generalization

Reader analysis

A learned task embedding tells the shared components which dynamics, reward, and action interface to use; masking prevents padded action dimensions from contaminating predictions or entropy. This architecture makes multitask learning possible, but its evaluation must still be split into distinct questions. The 80-task scaling experiment asks how well one model performs tasks represented in its offline data. The held-out-task experiment asks whether full-model online updates start from a useful initialization. Reader interpretation: neither alone establishes zero-shot control on arbitrary new embodiments. The appendix's hand-chosen task-embedding initialization and variable transfer curves sharpen that boundary. Likewise, changing from 15 to 30 to 80 tasks alters task composition as well as data. A clean claim about data scaling would need those factors controlled separately. e08e11e12e18e26

5.4 Training and inference

During training

Source description

Online learning alternates replay-based model updates and environment collection through planning. The encoder, dynamics, reward, and value components learn jointly; the next-observation embedding is stop-gradient. Policy-objective gradients update only the policy prior. No pretrained frozen backbone is specified. e04e05

Source description

Default training uses uniform replay, batch 256, update-to-data ratio 1, Adam, and temporal weighting 0.5. Table 8 gives joint-embedding/reward/value coefficients 20/0.1/0.1, which are omitted from the schematic Eq. (3). Reward and value heads have 101 bins in log-transformed space. e04e16e22

Source description

Offline scaling uses 545M transitions from 240 single-task agents for 80 tasks: 50 Meta-World and 30 DMControl. Behavior ranges from random to expert; batch size becomes 1024. The separate DMControl-only dataset has 345M transitions. e11

During inference

Source description

MPPI samples candidate action sequences, rolls them forward through latent dynamics, and ranks predicted discounted rewards plus terminal Q. It updates a diagonal Gaussian, executes its first action, observes the environment again, and replans. Policy-prior trajectories supplement sampling; the preceding plan initializes the next search after a one-step shift. e07

Source description

Table 8 specifies horizon 3, population 512, 24 policy-prior samples, 64 elites, and six planning iterations, increased by two when action dimensionality is at least 20. Thus the common recipe includes an explicit task-dependent rule. e22

5.5 Implementation flow

  1. Encode and predict in latent space

    The encoder h maps observation s and task embedding e to z. Dynamics d takes z, action a, and e to predict the next latent state. Reward R and terminal value Q evaluate actions; policy prior p proposes them. Future observations are never decoded. e03

  2. Stabilize the representation and targets

    LayerNorm and Mish operate in intermediate MLP layers. SimNorm applies softmax within latent groups, keeping each group on a simplex. Five Q-functions are used by default; TD targets take the minimum of two randomly sampled EMA target functions. e04e06e21e22

  3. Represent different tasks

    Every component receives a learned task embedding constrained to norm at most one. Inputs and outputs are padded to common dimensions; invalid action dimensions are excluded from policy predictions, entropy, and planning samples. e08

6. Experiments & results

TD-MPC2 learns a decoder-free latent world model and uses short-horizon planning, completed by a learned terminal value, to choose continuous actions. Its contribution is a coordinated robustness recipe and a task-conditioned architecture that supports larger multitask models; the evidence concerns simulated control, not generated-video quality or physical deployment.

6.1 Read the original evidence

Figure 4. Aggregate online learning favors TD-MPC2 across four simulation domains. Original paper, p. 5 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Read the horizontal axes as environment steps and compare methods within each domain. The first panel measures DMControl episode return; the other panels measure success rate on a percentage scale. Their vertical units therefore differ. The task counts printed above the panels sum to the paper's 104-task comparison. These are separately trained single-task agents, not one agent solving all four domains. The caption reports means and 95% confidence intervals over three seeds. Appendix C further specifies that success must hold at the final episode step, so briefly picking up and then dropping an object does not qualify. e09e10e17

What it supports. TD-MPC2's red aggregate curve is higher at the displayed endpoints in every domain, supporting the paper's broad robustness claim under its evaluation protocol. Aggregation does not imply a win on every task: the MyoSuite appendix specifically identifies Key Turn Hard as a case where TD-MPC succeeds earlier.

Where the evidence stops. The comparisons share environment protocols but differ in model size, update-to-data ratio, and tuning policy. These curves support data-efficiency comparisons under the reported settings; they do not establish superiority at equal training compute.

Table 1. Scaling improves the 80-task score, with diminishing gains per additional reported GPU day. Original paper, p. 7 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Read across a row to connect a model's size, approximate training cost, and final score. All rows concern the 80-task dataset, comprising 50 Meta-World and 30 DMControl tasks and 545M transitions collected from single-task agents. The caption specifies a single NVIDIA GeForce RTX 3090 as the cost reference. Score averages Meta-World success rates and DMControl returns normalized to the same 0–100 range. It is therefore an aggregate across different task metrics, not one overall success probability. The separate 30-task scaling curve and the held-out-task finetuning bars on this page belong to different experiments. e11e23

What it supports. The measured score rises from 16.0 at 1M parameters to 70.6 at 317M. As a reader calculation, the final size increase improves 68.0 to 70.6, or 2.6 score points, while reported cost rises from 12 to 33 GPU days. Larger models help, but the additional benefit becomes more expensive.

Where the evidence stops. The table provides no uncertainty intervals or collection-cost breakdown. Table 9 changes several architectural dimensions together, including ensemble size at the extremes. These observations support the tested configurations, not a universal scaling law.

Figure 17. Transfer improves the aggregate, while individual tasks reveal substantial variation. Original paper, p. 24 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Each panel is one task excluded from the 70-task pretraining set. Red denotes a pretrained 19M model updated online; gray denotes a similar-capacity model trained from scratch. Read the horizontal axis up to 40k environment steps. DMControl returns are normalized, while Meta-World uses success rate; both appear on a 0–100 scale. Bands are the reported 95% confidence intervals across three seeds. Appendix E specifies that the new task embedding is copied from a selected related task, such as Walker Walk for Walker Run. The complete model is finetuned, so this is not adaptation of only a small task-conditioning vector. e12e18

What it supports. Figure 8 summarizes a 47.0 versus 24.0 aggregate score at 20k steps, but this longer per-task view reveals the boundary. Pendulum Swingup and Hand Insert show clear advantages, while Hopper Hop remains near zero and Walker Run does not show a consistent gap. Transfer effectiveness depends on the target task.

Where the evidence stops. The source did not test random initializations or different source-task pairings for the new embedding. It therefore does not isolate how much benefit comes from pretrained dynamics, source-task selection, or their interaction.

6.2 Results and evaluation conditions

Task & protocolReported resultComparison & interpretation
Offline multitask control across 80 tasks

Fixed 545M-transition dataset; evaluation on the training task set, with increasing model configurations.

1M: 16.0 / 3.7; 5M: 49.5 / 4.2; 19M: 57.1 / 5.3; 48M: 68.0 / 12; 317M: 70.6 / 33.

Mean normalized score on a 0–100 scale; approximate GPU days

Each cost is reported for one NVIDIA GeForce RTX 3090; Figure 7 shows TD-MPC deteriorating as its size increases.

Larger TD-MPC2 configurations improve the aggregate score. This is neither unseen-task evaluation nor a fitted scaling law; uncertainty is not provided in Table 1. e11e23

Single-task online continuous control

104 tasks: 39 DMControl, 50 Meta-World, five ManiSkill2, ten MyoSuite; independent task training.

Figure 4 reports higher aggregate curves for TD-MPC2 across all four domains; means and 95% confidence intervals over three seeds.

Episode return for DMControl; final-step success rate for the other domains

SAC, DreamerV3, and TD-MPC; parameter counts, update ratios, and tuning policies differ.

Broad data-efficiency evidence, not a matched-compute comparison. Figure 16 identifies Key Turn Hard as an exception to TD-MPC2's per-task advantage. e09e10e17

Pick YCB object manipulation

Single-task ManiSkill2; all 74 YCB objects; 14M environment steps.

>60%, as stated in Figure 14's caption; curves show three-seed means and 95% confidence intervals.

Final-step success rate

SAC, DreamerV3, and TD-MPC fail to learn within the reported budget.

Successful simulated execution under this protocol; no real-robot transfer or held-out-object split is established. e10e17

Online adaptation to ten held-out tasks

19M model pretrained on 70 tasks, then full-model finetuning with empty replay on five DMControl and five Meta-World tasks.

47.0 finetuned versus 24.0 from scratch; three seeds.

Aggregate normalized score after 20k environment steps

Similar-capacity from-scratch agents; task embeddings initialized from selected semantically related source tasks.

An approximately twofold aggregate gain, not a universal per-task improvement. Alternative embedding initializations were not tested. e12e18

Planning and normalization ablations

Figure 9's 19M models trained on 80 tasks; separate online curves cover Dog Run, Humanoid Walk, and Pick YCB.

Policy 42.2; planning 53.7; planning plus policy 54.2. No normalization 46.8; SimNorm 51.0; LayerNorm plus SimNorm 54.2.

Normalized score

Within-figure actor and normalization variants.

Planning supplies most of the actor improvement. These ablation scores must remain distinct from the 19M scaling score of 57.1. e13e11

Test-time uncertainty regularization

Appendix J only: a fixed 19M, 80-task agent; no additional model updates.

No regularization 56.54; c=0.001: 58.14; c=0.01: 62.01; c=0.1: 44.13.

Normalized score

Same reported experiment with different planning-penalty strengths.

Moderate regularization helps and strong regularization hurts. It is absent from every other experiment; no uncertainty estimates are tabulated. e24

6.3 Ablations and diagnostic examples

Read component removals and qualitative examples within their stated evaluation conditions.

Figure 9. Planning and normalization explain substantial gains; the default is a practical combination of choices. Original paper, p. 8 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Each design choice has two views. Curves average online learning on Dog Run, Humanoid Walk, and Pick YCB; the adjacent bars evaluate 19M models across 80 tasks. Red marks the default TD-MPC2 formulation. Start with Actor: compare the policy alone, planning alone, and planning with policy proposals. Next examine how SimNorm and LayerNorm affect the normalization panels. The lower row changes the reward/value regression formulation or Q-ensemble size. Curve captions provide three-seed means and 95% confidence intervals. The bars start at 30, so their apparent heights exaggerate differences compared with a zero-based axis. e13e11

What it supports. Planning raises the actor score from 42.2 to 53.7, with policy proposals bringing it to 54.2. Normalization improves 46.8 to 51.0 with SimNorm and to 54.2 with LayerNorm added. Discrete regression scores 54.2 versus 49.6, while ten Q-functions score 57.0, above the default five-function score of 54.2.

Where the evidence stops. The multitask bars show no uncertainty intervals. Their default 54.2 differs from the 19M scaling score of 57.1 in Table 1; the source does not reconcile that difference. Keep each comparison within its reported experiment.

Table 10. An optional planning penalty helps at moderate strength and harms at high strength. Original paper, p. 30 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Read this as a test-time intervention on a 19M model trained on the 80-task dataset. The coefficient c changes how strongly planning penalizes uncertain trajectories; the model receives no additional gradient updates. Appendix J defines the penalty from the ensemble's mean value multiplied by its standard deviation and c, and subtracts it when ranking trajectories. Start with the unregularized column, then follow increasing c to see whether conservative search helps. The caption explicitly states that none of the paper's other experiments use this regularizer, so it is a separate extension rather than part of the default recipe. e24e11e13

What it supports. The score increases from 56.54 without regularization to 62.01 at c=0.01, a reader-calculated gain of 5.47 normalized-score points. Increasing c further to 0.1 reduces the score to 44.13. The experiment shows a useful range of penalty strengths for this model and dataset, not monotonic improvement from greater conservatism.

Where the evidence stops. No uncertainty intervals or per-task breakdown accompany the table. Its unregularized baseline also differs from earlier 19M results. Generalization of c=0.01 to other datasets is explicitly left for future work.

7. Analysis & limitations

7.1 What the evidence leaves open

Source description

The experiments are simulated continuous control. Discrete-action planning and broader zero-shot generalization remain open. The authors flag reward misspecification, unconstrained physical autonomy, and data cost as risks. e09e15

Source description

Visual RL uses a four-layer convolutional encoder, 64×64 images, and random-shift augmentation; results are comparable to leading baselines on ten DMControl tasks, not comprehensive visual generalization. e14

Reader analysis

Source inconsistencies remain: Eq. (4) says maximize a value term minus entropy while describing maximum-entropy learning; Appendix H reverses the temperature limits implied by softmax(z/τ). Figure 5/15 captions give 21 Humanoid actions, whereas Table 2 gives 24. e05e06e17e25

Reader analysis

Scaling changes layer widths, latent dimensions, encoder depth, and sometimes ensemble size. Appendix K also changes task composition with dataset size; it does not isolate a pure data-volume effect. e23e26

7.2 Questions for discussion

  1. How much of the planning advantage survives equalizing inference latency and model capacity?
  2. Would random or mismatched task-embedding initialization preserve the reported adaptation gain?

8. Reproducibility audit

8.1 Requirements and known gaps

Reader analysis

Reproduction needs the specified task versions and custom tasks, reward and final-step success logic, action repeats, mixed-quality replay, and Table 8/9 configurations. Appendix H gives useful settings but no complete version lock, reward/value bin support specification, or fixed-dataset optimizer-update budget. e10e11e22e23e25

Reader analysis

Before implementation, resolve the entropy-sign and temperature-description inconsistencies. Appendix H also calls T the number of tasks while using +T in input widths despite a 96-dimensional task embedding; actual tensor widths require clarification. e05e06e21e22

8.2 Proposed reproduction checks

The following checks are proposals motivated by the paper. They have not been run as part of this reading.

Check 1: Does SimNorm prevent instability under otherwise identical training?

Reader-proposed check, not performed: train the same 5M configuration on Walker Stand with SimNorm enabled or replaced by the identity, keeping LayerNorm, losses, Q ensemble, gradient clipping, replay sampling, and environment budget fixed. Use paired random seeds and the paper's four-million-step setting. Record episode return, latent norms, and both pre-clipping and post-clipping gradient norms throughout training. The hypothesis predicts that removing SimNorm increases pre-clipping instability and late performance collapse. Similar stable learning in both conditions would weaken a SimNorm-specific explanation, even if the complete TD-MPC2 recipe still outperforms TD-MPC. e06e20e21e22e27

Check 2: Does adaptation depend on the chosen task embedding?

Reader-proposed check, not performed: reuse a 19M model pretrained on the paper's 70-task split. For each of its ten held-out targets, compare the Table 7 source embedding with a random embedding and a deliberately mismatched source embedding. Keep pretrained weights, full-model updates, initially empty replay, budgets, and paired seeds identical; include the similar-capacity from-scratch control. Report per-task and aggregate scores at 20k and 40k steps with uncertainty. If random initialization preserves the gain, shared pretrained components carry much of the transfer benefit; if only the hand-selected pairing helps, the reported adaptation result depends materially on that initialization choice. e12e18

8.3 Reading coverage

Visual audit: Inspected the title and revision markings, all 22 figures and ten tables, method equations, evaluation definitions, and Appendix H architecture/hyperparameters. All six final original-PDF crops were individually viewed. The few-shot crop was expanded leftward and rendered at 300 DPI; inspection confirmed complete y-axis labels in both leftmost panels, all ten task panels, the rightmost 40k tick, and the full legend. Figure 3 was cross-checked against Eqs. (2)–(3) and the planning text: its omitted action inputs, task conditioning, and stop-gradient wiring are explained rather than inferred from the schematic. Source inconsistencies in entropy sign, SimNorm temperature limits, Humanoid dimensions, and architecture-width shorthand remain disclosed in the base report. Text reading covered all 31 pages through all eight chunks. Pages 10–15 contain related-work continuation, acknowledgements, references, and the appendix contents; these were read as text and supply no retained numerical or method claims. Separate supplements, external code, and videos remain outside this reading.

PDF pages inspected for this edition: 1, 2, 3, 4, 5, 6, 7, 8, 9, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31. Appendix coverage: reviewed.

Original figures and tables remain the work of the source’s authors. Extractions preserve their scientific content; any HTML wrapper layout is disclosed with each figure. The surrounding reading notes are our own.

Text reading scope & known omissions
  • Abstract; Sections 1–2: Introduction and Background
  • Section 3.1: Learning an implicit world model
  • Sections 3.2–3.3: MPC and generalist agents
  • Sections 4–4.1: Experiments and results
  • Sections 5–6: Lessons, opportunities, risks, and related work
  • Acknowledgements and references
  • Appendices A–C: Improvements, task visualizations, and task domains
  • Appendices D–G: Single-task, few-shot, additional ablation, and gradient results
  • Appendix H: Implementation details
  • Appendices I–K: Discrete actions, offline regularization, and additional multitask results

Outside the original text pass

  • Text extraction does not reconstruct figure images; inspect the retained PDF for figures and equation/table layout.
  • Separate supplemental material availability has not been fully verified.
  • The extraction limitation above was addressed by visually inspecting PDF pages 1–9 and 16–31, including every figure and table; pages 10–15 were read as text.
  • Identity: the inspected title page matches the catalog title and Nicklas Hansen; Hao Su; Xiaolong Wang. It identifies arXiv:2310.16828v2, 21 March 2024, published at ICLR 2024. No earlier revision or separate venue artifact was supplied, so revision-to-revision differences were not compared.
  • Code, linked resources, datasets, and checkpoints were not inspected; no experiments were reproduced.

The visual audit above records the subsequent illustrated pass.

8.4 Traceable evidence

e01PDF p. 1, title/author block and arXiv marginInspect

Title and authors match the catalog. All authors are affiliated with University of California San Diego; Hao Su and Xiaolong Wang are equal advisors. The artifact is arXiv:2310.16828v2, 21 March 2024, marked ICLR 2024.

Go to primary source ↓
e02PDF pp. 1–3, Sections 1–3Inspect

The motivation is robust learning from mixed-quality, multitask experience; the approach extends TD-MPC using implicit latent models and planning.

Go to primary source ↓
e03PDF p. 3, Section 3.1, Figure 3 and Eq. (2)Inspect

Five components encode observations, predict action-conditioned latent transitions, rewards and terminal values, and propose actions. Task embedding e conditions every component; observations are not decoded.

Go to primary source ↓
e04PDF p. 4, Section 3.1, Model objective, Eq. (3), and final architecture paragraphInspect

Joint model learning combines latent consistency with stop-gradient targets and discrete reward/value regression. Q targets use EMA functions and the minimum of two sampled members from a default five-member ensemble.

Go to primary source ↓
e05PDF p. 4, Policy objective, Eq. (4) and explanatory paragraphsInspect

The prose describes a stochastic maximum-entropy policy and gradients only through p. The displayed maximization objective subtracts β times H, explicitly called entropy, leaving a sign inconsistency.

Go to primary source ↓
e06PDF p. 4, Eq. (5); pp. 28–29, Appendix H, Simplicial NormalizationInspect

SimNorm partitions the latent vector and applies softmax within each partition. Eq. (5) divides logits by temperature; p. 29's claims about infinite and zero temperature reverse the corresponding mathematical limits. The displayed implementation uses ordinary groupwise softmax.

Go to primary source ↓
e07PDF pp. 4–5, Section 3.2, Eq. (6) and continuationInspect

MPPI searches Gaussian action sequences using predicted rewards and terminal Q, mixes in policy-prior trajectories, warm-starts from the shifted previous solution, executes the first action, and repeats.

Go to primary source ↓
e08PDF p. 5, Section 3.3, Learnable task embeddings and Action maskingInspect

Task embeddings are learned jointly and constrained to norm at most one. Zero-padding and action masking accommodate different dimensions during training and planning.

Go to primary source ↓
e09PDF pp. 5–6, Figure 4, Section 4 and BaselinesInspect

The 104-task online comparison covers four simulation domains. Figure 4 gives aggregate means and 95% confidence intervals over three seeds. Default TD-MPC2 has 5M parameters and UTD 1; DreamerV3 uses about 20M parameters and reported UTD 512.

Go to primary source ↓
e10PDF p. 20, Appendix C, Environment details and Table 6Inspect

Episodes have fixed lengths with no termination conditions; success is assessed at the final step. The table specifies per-domain action repeat, effective episode length, training steps, and evaluation metric.

Go to primary source ↓
e11PDF p. 7, Section 4.1, Figure 7 and Table 1Inspect

The 80-task dataset has 545M transitions from 240 agents; its 30-task DMControl subset has 345M. Batch is 1024. Table 1 pairs 1/5/19/48/317M models with scores 16.0/49.5/57.1/68.0/70.6 and approximately 3.7/4.2/5.3/12/33 GPU days on one RTX 3090. Score averages success rates and normalized returns.

Go to primary source ↓
e12PDF pp. 7–8, Figure 8 and Few-shot learningInspect

A 19M model pretrained on 70 tasks is finetuned online on ten held-out tasks with initially empty replay and full-model updates. At 20k steps the aggregate is 47.0 versus 24.0 from scratch; three seeds.

Go to primary source ↓
e13PDF p. 8, Figure 9, all actor, normalization, regression and Q-function panelsInspect

The actor bars are 42.2/53.7/54.2; normalization 46.8/51.0/54.2; continuous/discrete regression 49.6/54.2; two/five/ten Q-functions 53.5/54.2/57.0. Online curves average three hard tasks and three seeds; multitask bars concern 19M agents on 80 tasks.

Go to primary source ↓
e14PDF pp. 8–9, Figure 10 and Visual RL; p. 29, Appendix H, Visual RLInspect

Ten image-based DMControl tasks use a four-layer convolutional encoder with 64×64 inputs and random-shift augmentation. Performance is described as comparable to DrQ-v2 and DreamerV3; curves report three-seed means and 95% confidence intervals.

Go to primary source ↓
e15PDF p. 9, Section 5; p. 29, Appendix IInspect

The authors identify discrete actions and broader generalization as future directions and discuss reward, physical-autonomy, and data-access risks.

Go to primary source ↓
e16PDF p. 16, Appendix A, Summary of ImprovementsInspect

Changes include LayerNorm/Mish, SimNorm, Q-ensemble dropout, maximum-entropy prior, discrete regression, uniform replay, removal of MPPI momentum, and multitask conditioning.

Go to primary source ↓
e17PDF p. 23, Figures 14–16 and captions; p. 6, Figure 5 captionInspect

Figure 14 reports >60% Pick YCB success after 14M steps on all 74 objects and baseline failure within that budget. Figure 16 identifies Key Turn Hard as an exception. Figures 5 and 15 label Humanoid as 21-dimensional.

Go to primary source ↓
e18PDF p. 24, Appendix E, Figure 17 and Table 7Inspect

Finetuning efficacy varies by task. The ten targets use named source-task embeddings, including Walker Run from Walker Walk and Door Lock from Door Open. Alternative initializations or pairings were not tested.

Go to primary source ↓
e19PDF p. 25, Appendix F, Figures 18–20Inspect

Normalized task embeddings score 54.2 versus 46.6 without normalization. T-SNE is qualitative; activation comparisons show similar performance but smoother Mish gradients.

Go to primary source ↓
e20PDF p. 26, Appendix G, Figure 21Inspect

Gradient norms are plotted on a log scale for five DMControl tasks, one displayed seed per task. TD-MPC2 gradients remain comparatively stable while TD-MPC often grows sharply.

Go to primary source ↓
e21PDF p. 26, Appendix H, Architectural details and displayed architectureInspect

The source shows separate encoder, dynamics, reward, policy, Q ensemble and task embeddings, with 96-dimensional embeddings. Its width notation uses +T while defining T as task count, leaving the dimensional shorthand inconsistent.

Go to primary source ↓
e22PDF p. 27, Appendix H, Table 8Inspect

Defaults include horizon 3; 6 planning iterations plus 2 for action dimensions at least 20; population 512; 24 policy samples; 64 elites; uniform replay; batch 256; UTD 1; 512-dimensional latent state; SimNorm groups of 8 at temperature 1; five Q-functions; 101 bins; Adam; loss coefficients 20/0.1/0.1 and temporal coefficient 0.5.

Go to primary source ↓
e23PDF p. 28, Appendix H, Eqs. (7)–(8) and Table 9Inspect

Discount and seed steps depend on effective episode length. Scaling changes widths, latent dimension, encoder depth, and Q-ensemble size; the largest configuration uses eight Q-functions. No other model configurations were tested.

Go to primary source ↓
e24PDF p. 30, Appendix J, Eqs. (9)–(12) and Table 10Inspect

A planning-only uncertainty penalty is c times ensemble mean times ensemble standard deviation. Scores are 56.54 without regularization, 58.14 at 0.001, 62.01 at 0.01, and 44.13 at 0.1. The paper explicitly excludes this regularizer from all other experiments.

Go to primary source ↓
e25PDF pp. 18–19, Appendix C, Tables 2–5; p. 29, Appendix H, baseline descriptionsInspect

Task tables list dimensions, including 24 for Humanoid and seven for Pick YCB. Table 3 abbreviates the Meta-World list with an ellipsis. Appendix H names baseline implementations and settings without a complete software version lock.

Go to primary source ↓
e26PDF p. 31, Appendix K, Figure 22Inspect

A random 15-task DMControl subset is compared with the 30- and 80-task datasets. Scores improve with model size; different task compositions prevent interpreting the comparison as a controlled data-volume experiment.

Go to primary source ↓
e27PDF p. 21, Appendix D, Figure 12, Walker Stand panel and captionInspect

Walker Stand is shown over four million environment steps. TD-MPC's late return deterioration contrasts with TD-MPC2's stable return; Appendix G provides the related gradient diagnostic.

Go to primary source ↓

8.5 Primary sources

Scroll across the image to inspect details. Press Esc to close.