Hierarchical World Models as Visual Whole-Body Humanoid Controllers

In plain words

Teaching a simulated humanoid to walk, run and jump over gaps from camera images is hard: a human body has many joints, and falling over is easy. Earlier methods often simplify the problem: they use low-dimensional privileged observations, pre-built skills, or train each task from scratch.

This paper takes a different route. First, one general controller learns to copy many recorded human motions. It learns the body's physics, not any single task. Second, a separate "puppeteer" learns each new task from pictures. It does not move the joints itself. It only tells the frozen controller where the hands, feet and head should go next. Both parts are learned by planning with a model of the world: each one imagines possible futures and picks the best.

The result is a body that moves in a human-looking way, and the controller can be reused across tasks. The authors report that people preferred its motions in a 51-person study.

TL;DR

• Two TD-MPC2 world models are trained separately: a low-level tracker that follows MoCap-derived end-effector targets, and a visual high-level planner that outputs those targets §3.
• The tracker is pretrained once on 836 clips and frozen; the planner is trained from scratch per task on task reward §3.1 §3.2 §4.1.
• On 8 humanoid tasks it matches TD-MPC2 return except on stairs; motions are preferred in a 51-person study (author-reported) §4.2.
• Relevant if you need a whole-body humanoid controller with human-like motion from pixels and cannot write skill primitives. Code is MIT but the setup is heavy (about 12 days of tracker training on one RTX 3090) §4.1.

Problem

Whole-body humanoid control with a high-dimensional action space (56 joints) and an unstable bipedal body is hard for reinforcement learning. Prior work either uses privileged low-dimensional observations, learns skill primitives, or trains each downstream task from scratch with adversarial motion rewards. Some prior work learns up to about 2,600 separate low-level tracking policies §1 §3. The authors want a data-driven, visual method with few assumptions, no reward shaping and no skill primitives, that also yields natural motion §1.

Method

The method is two hierarchical agents, both TD-MPC2 world models trained independently §3. Each model has the same six components, written as latent state encoder, latent dynamics, reward, termination, terminal value and policy prior eq. 1. The authors add a termination head D(z, a) that predicts the probability that the episode ends eq. 1 §3.3. Their description of the planner says actions are chosen by MPC with MPPI, with the policy prior used as a warm start §2.

Low-level tracking agent §3.1. Input is the proprioceptive state q_t and a command c_t. The command is 15 numbers: relative 3D positions of 5 end-effectors (head, two hands, two feet) taken from a reference motion §3.1 app. E. The output is a sequence of H actions that approximately follows the command. Training data are sampled sub-sequences from MoCapAct, an offline set of noisy rollouts from expert tracking policies for 836 MoCap clips §3.1 app. D. During training the command is the reference motion's future end-effector positions §3.1. Each reward is labelled with the CoMic tracking reward §3.1. Offline and online data are mixed 50/50 in each gradient batch §3.1 §4.1. One model is trained for all 836 clips §4.1.

High-level puppeteering agent §3.2. Input is the same proprioceptive state plus a 64x64 RGB image from a third-person camera §4.1 app. E. The agent plans over commands of the same shape as the tracker's input, i.e. the action space is the command space §3.2. The tracker stays frozen §3.2. Both agents act at the same frequency, which the authors write as k = 1 §4.1. The puppeteer is trained from scratch per task, using only the downstream reward §3.2 §4.1.

Termination handling §3.3. Episodes end on timeout (500 steps for the downstream tasks) or when a non-foot part touches the floor §4.1. Two changes follow. First, the termination head is trained with cross-entropy against the true termination signal, added to the TD-MPC2 loss with coefficient alpha eq. 2. Second, TD targets are truncated at terminal states §3.3. At test time the predicted termination is used as a soft truncation: a running discount of the predicted termination probabilities down-weights rollout values §3.3.

TD-MPC2 losses (not written out in this paper). The paper refers to the TD-MPC2 objective without restating it eq. 2 §2. The components are latent consistency, reward, and value terms, which the code weights with coefficients from tab. 5. The details below come from the code, not the paper:

# One world-model update on sampled sub-sequences  [§2] [eq. 2]
# Input:  obs (H+1, B, ...), action (H, B, ...), reward (H, B),
#         terminated (B,)
# Output: scalar loss

def update(obs, action, reward, terminated):
    # Encode the next observations as stop-gradient targets  [§2]
    next_latents = encoder(obs[1:])
    # TD target from target Q-nets: min of two random heads
    # [code: tdmpc2.py#L204]
    td_targets = reward + gamma * (1 - terminated) * q_target(next_latents)
    # Roll the latent model forward H steps from the first observation  [§2]
    latents = [encoder(obs[0])]
    for t in range(horizon):
        latents.append(dynamics(latents[-1], action[t]))
    # Consistency: predicted latents match encoded next observations  [§2]
    consistency = mse(latents[1:], next_latents)
    # Reward and value heads: cross-entropy on two-hot targets  [§2]
    reward_loss = soft_ce(reward_head(latents, action), reward)
    value_loss = soft_ce(q_heads(latents, action), td_targets)
    # Termination: BCE of the final latent against the last step's flag  [eq. 2]
    term_loss = bce(terminated_head(latents[-1]), terminated)
    # Weight time steps by 0.5 to the power t  [tab. 5]
    # [code: tdmpc2.py#L234]
    total = (20 * consistency + 0.1 * reward_loss
             + 0.1 * value_loss + 0.1 * term_loss)
    # Update the model, then the policy prior, then the target Q-nets  [§2]
    return total

# Predicted return of one action sequence with soft termination  [§3.3]
# Input:  latent (N, 512), actions (H, N, 15)
# Output: value (N,)

def estimate_value(latent, actions):
    # Accumulate the reward of each step, weighted by survival  [§3.3]
    total = 0
    discount = 1
    terminated = 0
    for t in range(horizon):
        total += discount * (1 - terminated) * reward(latent, actions[t])
        latent = dynamics(latent, actions[t])
        discount = discount * gamma
        # Running termination probability, capped at 1  [code: tdmpc2.py#L98]
        terminated = min(terminated + terminate_prob(latent), 1)
    # Add the terminal value of the policy's action at the last latent  [§2]
    terminal = q_avg(latent, policy_prior(latent))
    return total + discount * (1 - terminated) * terminal

# MPPI planning with a warm-started, clipped Gaussian
# [§2] [code: tdmpc2.py#L103]
# Input:  latent (512,)
# Output: first action (A,) with A = 15 (puppeteer) or 56 (tracker)

def plan(latent):
    # Seed 24 candidate sequences from the policy prior  [tab. 5]
    prior_actions = rollout_policy_prior(latent, samples=24, horizon=3)
    # Start from the previous plan's mean, shifted by one step;
    # zeros at the first step of an episode  [code: tdmpc2.py#L127-L129]
    mean = shift(previous_mean)
    # Initial spread is max_std = 2  [code: config.yaml#L44]
    std = full((horizon, A), 2.0)
    for iteration in range(iterations):
        # Sample 512 - 24 sequences, clipped to [-1, 1]  [§2] [tab. 5]
        actions = clip(mean + std * noise, -1, 1)
        # Score all sequences  [§2]
        values = estimate_value(latent, actions)
        # Keep the 64 best (elites)  [tab. 5]
        elites = top_k(values, 64)
        # Refit mean and std with softmax weights at temperature 0.5  [tab. 5]
        weights = softmax(0.5 * elites.values)
        mean = weighted_mean(weights, elites.actions)
        std = clip(weighted_std(weights, elites.actions), 0.05, 2.0)
    # Draw the first action from the elites; in training add noise with the
    # first step's std  [code: tdmpc2.py#L152-L161]
    action = sample_first_action(weights, elites)
    return clip(action + std[0] * noise, -1, 1)

Hierarchy at execution time, as implemented in the official code:

# One control step of the two-level hierarchy  [§3] [§3.2]
# Input:  env state, camera image; frozen tracker
# Output: env observation after the step

def hierarchy_step(state, image):
    # Puppeteer plans a 15-value command in [-1, 1] from state and image  [§3.2]
    command = puppeteer.act(state, image)
    # Code scales the command by 0.3 and adds it to the current end-effector
    # positions, so commands are offsets, not absolute targets
    # [code: puppeteer/envs/transfer.py#L120-L130]
    targets = end_effector_positions(state) + 0.3 * command
    # Frozen tracker maps state and targets to 56 joint actions, k = 1
    # [§3.2] [§4.1]
    joint_action = tracker.act(state, targets)
    return env.step(joint_action)

Hyperparameters

Values from tab. 5 unless noted. "Code" entries show where the official repository differs or adds a value that the paper omits. The same hyperparameters are used for all tasks and both levels, except where noted tab. 5.

Compute §4.1 tab. 3: one RTX 3090. The tracker takes about 12 days and the puppeteer about 4 days of training. The paper's table gives wall time of 21.8 hours per million steps for Puppeteer and 29.0 hours with vision, and 88.2 ms inference per step for Puppeteer versus 50.8 ms for TD-MPC2 tab. 3.

Results

Author-reported. Setup: 8 tasks on a simulated 56-DoF CMU humanoid, each run with 10 random seeds and a 3M environment-step budget §4.2 fig. 5. Baselines are SAC, DreamerV3, TD-MPC2 (with contact termination added), and two hierarchical variants that replace the puppeteer with SAC or DreamerV3 §4.1.

Notes on the table:

• The human study covers only the five visual tasks. In each of the 10 pairs, Puppeteer was preferred in 47 to 51 of 51 responses tab. 6.
• Under "Results" the paper says that 4% of responses rated the methods equally natural and 0% preferred TD-MPC2 §4.2. Table 6 gives 1.8% and 0.4% tab. 6. Both are reported here, and the table is the source for the counts.
• Tracker data rows tab. 2 app. B are monotonic: offline-only data is the worst on all three metrics, and each larger share of clip data improves them up to the full mix.
• Ablations fig. 8 (mean over 3 seeds for pretraining; downstream ablations averaged over 5 tasks): the text says planning at both levels is critical. Pretraining the high-level agent is described as beneficial, but this is only suggested §4.3. Finetuning from a corridor-pretrained model is described as helping substantially on the other visual tasks §4.3.
• No per-task return values are given in text or tables; only the curves in fig. 5 show them.

Paper vs code

Official code: https://github.com/nicklashansen/puppeteer, linked from the first page and from the project page. Commit f9e3ed7ad4bf97380bd35c158469d1e5689438d9, MIT license. The repository owner is the first author (Nicklas Hansen), so it belongs to an author. The commit date was not checked. Model checkpoints are linked from the README, not from the paper.

Discrepancies found, paper first, then code:

• Termination head input. Paper: D(z, a), a function of latent and action eq. 1. Code: terminated(z), latent only world_model.py world_model.py.
• Termination loss. Paper: cross-entropy between predicted and true termination signals eq. 2; the paper does not say over which time steps. Code: a single binary cross-entropy on the final latent of each sampled sub-sequence, against the flag of its last step only tdmpc2.py tdmpc2.py buffer.py.
• Corridor reward. Paper: reward proportional to forward velocity, clipped to [0, v_target] with v_target = 6 eq. 3. Code: a tolerance centred on v = 6 with linear falloff and margin v, so it also falls for speeds above 6 run_through_corridor.py.
• Proprioceptive reward. Paper: min of planar speed and target plus alpha times head height, with alpha not stated eq. 4. Code: product of a standing term (height constant 1.6), a small-control term, a hand-position term and a movement term; walk speed 1.0 and run speed 6.0 walk.py walk.py. The code's reward is not the paper's formula.
• Termination condition. Paper: non-foot contact with the floor §4.1. Code adds: any end effector below z = -0.5 also ends the episode run_through_corridor.py.
• Command semantics. Paper: the puppeteer outputs geometric locations for end effectors §1 §3. Code: the command is scaled by 0.3 and added to the current end-effector positions transfer.py config.yaml.
• Tracker reference window. Paper: the command holds reference positions for t+1 through t+H §3.1. Code: one future frame (ref_steps=[1]), giving 15 values tracking.py.
• Planner iterations. Paper: 8 tab. 5. Code: 6 in config, plus 2 only for action dimension 20 or more tdmpc2.py. The 15-D puppeteer gets 6 iterations.
• Seed steps. Paper: 2,500 tab. 5. Code: max(1000, 5 times episode length) [code: puppeteer/envs/init.py#L33], i.e. 5,000 for the 1000-step tracker. At the end of seeding the code runs as many updates as there were seed steps, in one burst trainer.py.
• Encoder learning rate. Paper: 1e-4 tab. 5. Code: 3e-4 times 0.3 = 9e-5 tdmpc2.py config.yaml.
• Latent normalisation. Paper: LayerNorm + Mish tab. 5. Code: SimNorm on encoder and dynamics outputs (group size 8) layers.py world_model.py.
• Image preprocessing. Paper: raw 64x64 RGB with no preprocessing app. E. Code: random-shift augmentation (pad 3) and scaling to [-0.5, 0.5] layers.py.
• Success metric. Paper appendix: tracking success means error stays under 0.5 for up to 100 steps app. B. Code: the episode's success flag is set when the clip is reached humanoid.py. No 0.5 threshold was found in the repository, so the tracker numbers may not match the code's success flag.
• Training budget. Paper: 3M environment steps for downstream tasks §4.2. Code default: 10M steps for every task config.yaml.
• Camera. Paper: third-person camera §4.1. Code: camera 0 for corridor tasks and 3 for the others transfer.py. Which camera is third-person was not checked in the MJCF assets.
• Dropout and tau. Paper: neither stated. Code: Q dropout 0.01 and target update 0.01 config.yaml config.yaml.

Pitfalls

• Termination is the hard part. The paper calls termination in planning "novel" §3.3. The code trains the head on the last step only and applies a clamped running sum at plan time, so expect noisy values and check the sum is capped at 1 tdmpc2.py.
• Commands are relative offsets, not absolute targets. Read the official code before building a puppeteer. The paper's wording suggests absolute end-effector locations §1 §3.
• Reward for corridor tasks differs from the equation. Implement the code's tolerance-based reward if you want to match the reported results; the paper's clipped linear form would reward faster-than-target runs equally eq. 3 run_through_corridor.py.
• TD-MPC2 losses are not written out. The consistency, reward and value terms, their two-hot / symlog encoding (101 bins over [-10, 10]), the min over two random Q heads, and the tanh squashing of the policy are all in the code only math.py tdmpc2.py.
• Dimension-dependent planner. The +2 iterations for large action spaces rule is a heuristic in the code tdmpc2.py. Set iterations explicitly if you want the paper's 8 iterations for both levels.
• Planning cost. Each action is planned with 512 samples, 8 iterations and horizon 3 (6 iterations for the puppeteer). Inference takes 88.2 ms per step for Puppeteer against 50.8 ms for TD-MPC2 tab. 3.
• Pretraining data is third-party. The tracker uses MoCapAct's rollouts of expert policies, which depend on the CMU MoCap data app. D §3.1. Check licences of those sources before reuse; the paper states MoCapAct is MIT-licensed § reproducibility statement.
• Code default seed burst. At the end of seeding the trainer runs one update per seed step at once, which the paper does not describe trainer.py.
• Figure and table numbers were counted from the LaTeX source in the order the captions appear. Check them against the compiled PDF before quoting in print.
• Compute and length. Training time quoted in text (about 12 days for the tracker) does not match 21.8 hours per million steps times 10 million steps (about 9 days) in tab. 3. Treat both as approximate.
• Seeds. The benchmark curves use 10 seeds fig. 5; the ablations use 3 seeds fig. 8 fig. 9.

Lineage

• Builds on TD-MPC2 (Hansen et al., ICLR 2024), the world-model backbone of both agents §2. No card for it exists in this wiki yet.
• Builds on MoCapAct (Wagener et al., NeurIPS 2022), whose offline rollouts and CMU humanoid setup are used for the tracker §3.1 app. D.
• Uses the CoMic tracking reward (Hasenclever et al., 2020) for labels §3.1.
• Compared against SAC, DreamerV3, TD-MPC2 and hierarchical variants with SAC or DreamerV3 at the high level §4.1.
• Related work on skill-based and motion-prior controllers §5. Concurrent work HumanoidBench uses a different embodiment and only privileged state tasks §5.
• Superseded by: not known from content on this wiki.