Long-horizon robot planning demands predictive representations at different timescales and levels of abstraction. While high-frequency motor control requires tracking millimeter-level contact geometries and instantaneous joint velocities at 50 Hz, high-level task planning requires reasoning over macro-milestones spanning seconds or minutes.
Traditional pixel-generative world models try to predict raw visual observations across extended futures. This strategy carries severe drawbacks:
The Joint-Embedding Predictive Architecture (JEPA), pioneered by LeCun and Balestriero, bypasses pixel generation entirely. Instead of predicting raw observation frames $\hat{o}_{t+1} \in \mathbb{R}^{H \times W \times C}$, a JEPA optimizes predictions purely in a learned latent representation space:
\[\hat{z}_{t+1} = \text{Pred}_\theta(z_t, a_t), \qquad z_{t+1} = \text{Enc}_\phi(o_{t+1}) .\]Hierarchical world models scale this paradigm to multi-timescale predictive abstraction. The architecture trains a multi-level pyramid of latent world models end-to-end. Each level operates over a distinct temporal stride and predicts forward in its own learned latent space:
```lab hjepalab {“dock”: true, “title”: “7-DOF Franka Panda Multimodal Sensor Lab”, “height”: 420}
Planning runs top-down: the highest level optimizes macro-actions toward the target goal, its predictions serve as intermediate subgoals for the intermediate level, and the lowest level unrolls concrete motor primitives for closed-loop execution.
## Hierarchical predictive learning
Hierarchical world models learn this multi-timescale latent dynamics directly from observation-action trajectories without generative decoders or hand-engineered state abstractions.
<div class="hjepa-interactive-frame">
<iframe src="media/hjepa/interactive/hierarchical-training-interleaved.html?embed=1" style="width:100%; height:620px; border:none;" title="Interactive hierarchical training animation" loading="lazy"></iframe>
</div>
At each level $l \in \{1, \dots, L\}$, the architecture consists of four modular components:
1. **Observation Encoder** $\text{Enc}^{(l)}$: Maps lower-level representations into level-$l$ latent vectors. Level 1 processes raw multimodal observations (RGB camera frames, depth maps, proprioceptive joint states). Level 2 processes a temporal window of Level 1 latent embeddings $z_{t : t + W}^{(1)}$.
2. **Action Encoder** $\text{Enc}_a^{(l)}$: Compresses a sequence of lower-level actions over temporal stride $\sigma_l$ into a macro-action embedding $a^{(l)}$. At Level 1, this represents raw motor torques or Cartesian velocity vectors.
3. **Action-Conditioned Predictor** $\text{Pred}^{(l)}$: An autoregressive Transformer or MLP that advances the latent state conditioned on macro-action $a_t^{(l)}$:
$$\hat{z}_{t+1}^{(l)} = \text{Pred}^{(l)}(z_t^{(l)}, a_t^{(l)}) .$$
4. **Regularization Objective**: Self-supervised non-contrastive regularization via SIGReg (Skew-Independent Gaussian Regularization) to prevent representation collapse without requiring negative pairs or heuristic momentum encoders.
Gradients from higher-level prediction losses flow directly through lower-level encoders, training the entire hierarchy jointly end-to-end.
## Hierarchical planning from coarse to fine
Planning in hierarchical world models optimizes actions across the temporal hierarchy in a top-down fashion. Instead of running expensive sample-based planning (such as the Cross-Entropy Method or Model Predictive Path Integral control) over thousands of primitive motor steps, the search space is factored hierarchically:
<div class="hjepa-interactive-frame">
<iframe src="media/hjepa/interactive/hierarchical-planning.html?embed=1" style="width:100%; height:620px; border:none;" title="Interactive hierarchical planning animation" loading="lazy"></iframe>
</div>
1. **Top-Level Optimization**: Given current observation $o_0$ and target goal $o_g$, the highest level computes latent representations $z_0^{(L)} = \text{Enc}^{(L)}(o_0)$ and $z_g^{(L)} = \text{Enc}^{(L)}(o_g)$. The planner optimizes a sequence of macro-actions to minimize terminal latent distance:
$$\min_{\{a_t^{(L)}\}} \mathcal{D}\Big( \hat{z}_{H_L}^{(L)}, z_g^{(L)} \Big) .$$
2. **Subgoal Propagation**: The unrolled trajectory states $\hat{z}_1^{(L)}, \hat{z}_2^{(L)}, \dots$ serve as intermediate goal targets for Level $L-1$.
3. **Low-Level Motor Synthesis**: The base level unrolls continuous action primitives $\hat{a}_{0:K}^{(1)}$ to reach each short-horizon waypoint, executing actions on the physical robot in a closed feedback loop.
This temporal factoring reduces the effective search horizon from hundreds of raw timesteps to a small sequence of structured latent decisions, dramatically reducing planner compute while escaping local minima.
## Mathematics
### Joint-embedding predictive energy
Let an observation trajectory be denoted by $(o_1, a_1, o_2, a_2, \dots, o_T)$. For hierarchy level $l$, the forward prediction error measures alignment between predicted latent state and encoded future observation:
$$\mathcal{L}_{\text{pred}}^{(l)}(\theta, \phi) = \frac{1}{T_l - 1} \sum_{t=1}^{T_l - 1} \Big\| \text{Pred}_\theta^{(l)}\big(z_t^{(l)}, a_t^{(l)}\big) - z_{t+1}^{(l)} \Big\|_2^2 ,$$
where $z_t^{(l)} = \text{Enc}_\phi^{(l)}(s_t^{(l-1)})$ is the target embedding produced by the encoder.
### SIGReg: Collapse prevention without heuristics
In non-generative joint embeddings, minimizing predictive loss alone causes dimensional collapse: the encoder can trivially project all observations to a constant vector $z_t = \mathbf{c}$, reducing prediction error to zero.
The architecture utilizes SIGReg (Skew-Independent Gaussian Regularization). Rather than using data augmentation pairs or momentum target networks, SIGReg forces the distribution of encoded embeddings across the batch $Z = [z_1, \dots, z_B]^\top \in \mathbb{R}^{B \times d}$ to match an isotropic standard normal distribution $\mathcal{N}(0, I_d)$.
The regularizer decomposes into two penalty terms:
$$\mathcal{L}_{\text{SIGReg}}(Z) = \lambda_{\mu} \|\bar{z}\|_2^2 + \lambda_{\Sigma} \Big\| \frac{1}{B-1} (Z - \bar{z})^\top (Z - \bar{z}) - I_d \Big\|_F^2 ,$$
where $\bar{z} = \frac{1}{B} \sum_{i=1}^B z_i$ is the empirical batch mean, and $\|\cdot\|_F$ denotes the Frobenius norm. By constraining the empirical covariance to the identity matrix $I_d$, all latent dimensions remain mutually uncorrelated and non-degenerate.
### Inverse Dynamics Regularization for heterogeneous robot data
On multi-scene robotic datasets such as DROID, camera viewpoints, lighting, and table backgrounds vary drastically across demonstrations. A purely predictive objective can minimize loss by modeling the static scene background while ignoring small manipulator motions.
To guarantee that latent states retain control-relevant information about the robot's physical configuration, the architecture incorporates an Inverse Dynamics Model (IDM) objective:
$$\mathcal{L}_{\text{IDM}}^{(l)}(\psi) = \frac{1}{T_l - 1} \sum_{t=1}^{T_l - 1} \Big\| \text{IDM}_\psi^{(l)}\big(z_t^{(l)}, z_{t+1}^{(l)}\big) - a_t^{(l)} \Big\|_2^2 .$$
<div class="hjepa-interactive-frame">
<iframe src="media/hjepa/interactive/idm-schema.html?embed=1" style="width:100%; height:620px; border:none;" title="Interactive Inverse Dynamics animation" loading="lazy"></iframe>
</div>
The complete joint training objective optimizes:
$$\mathcal{L}_{\text{total}} = \sum_{l=1}^L \Big( \mathcal{L}_{\text{pred}}^{(l)} + \alpha_l \mathcal{L}_{\text{SIGReg}}^{(l)} + \beta_l \mathcal{L}_{\text{IDM}}^{(l)} \Big) .$$
### Multimodal physical sensor fusion: RGB-D, LiDAR, Radar, and Raman
Physical robot perception operates across heterogeneous sensory bands beyond optical RGB cameras. When manipulating objects in industrial, chemical, or unstructured environments, optical vision alone is degraded by occlusions, specular reflections, ambient dust, and visually indistinguishable materials.
Hierarchical world models natively accommodate multimodal physical sensing by mapping distinct sensory streams into token embeddings before latent projection:
1. **3D RGB-D Vision**: High-resolution photometric pixels aligned with metric depth maps ($0.20$ to $2.50\,\text{m}$), providing global scene context and visual edge geometry: $e_{\text{RGB-D}} \in \mathbb{R}^{d}$.
2. **3D Time-of-Flight LiDAR**: 64-beam pulsed laser range returns measuring calibrated metric clearances and surface reflectivity intensity: $e_{\text{LiDAR}} \in \mathbb{R}^{d}$.
3. **77 GHz mmWave FMCW Radar**: Range-Doppler velocity spectrograms measuring radial velocity $v_r = \frac{c \Delta f_d}{2 f_0}$ and radar cross-section ($\text{RCS}$ in $\text{dBm}^2$), penetrating dust, smoke, and visual clutter to track moving bodies: $e_{\text{Radar}} \in \mathbb{R}^{d}$.
4. **Wrist-Mounted Raman Spectroscopy**: Inelastic molecular laser scattering ($\lambda = 785\,\text{nm}$) measuring vibrational intensity across wavenumbers ($400 - 3200\,\text{cm}^{-1}$). Identifies chemical composition (distinguishing polyethylene polymers from nylon polyamides and aluminum alloys) for material-grounded manipulation: $e_{\text{Raman}} \in \mathbb{R}^{d}$.
5. **Proprioception & Joint Telemetry**: 7-DOF joint angles $q$, joint velocities $\dot{q}$, and torque commands $\tau$: $e_{\text{proprio}} \in \mathbb{R}^{d}$.
The level-1 observation representation combines these modal tokens into a unified embedding:
$$z_t^{(1)} = \text{Enc}_{\text{fusion}}\big( [e_{\text{RGB-D}}, e_{\text{LiDAR}}, e_{\text{Radar}}, e_{\text{Raman}}, e_{\text{proprio}}] \big) .$$
#### The non-generative JEPA advantage
Generative world models fail when scaling to multimodal physical suites. Synthesizing raw radar chirp spectrograms, ToF range distributions, chemical Raman spectra, and pixel grids simultaneously requires modeling complex, multimodal noise characteristics across incompatible physics engines.
Because hierarchical world models predict strictly within learned latent embeddings without generative decoders, they fuse complementary physical sensors without generating raw sensor outputs. SIGReg regularizes the fused representation, preventing dominant modalities (such as dense pixels) from suppressing sparse modalities (such as Raman spectra or radar Doppler velocity).
### 3D in-scene sensor rendering during robot rollout
In the interactive lab above, physical sensor signals are rendered directly within the three-dimensional MuJoCo simulation viewport during execution:
1. **Raman Spectroscopic Laser Beam & Reticle**: A ruby red laser beam ($\lambda = 785\,\text{nm}$) projects from the Franka tool center point down to the target workpiece surface, with a circular reticle marking the laser excitation footprint.
2. **3D Time-of-Flight LiDAR Rays**: A 32-beam radial fan of pulsed cyan laser scan lines radiates from the manipulator toward the workcell, terminating in point cloud surface hit markers on obstacles, receptacles, and parts.
3. **77 GHz mmWave FMCW Radar Cone**: A translucent amber electromagnetic wave cone projects from the sensor toward moving targets, measuring Doppler frequency shifts to track workpiece momentum and velocity.
4. **RGB-D 3D Voxel Bounding Boxes**: Real-time wireframe bounding boxes enclose tracked parts, indicating spatial bounding dimensions and estimated centroids.
### Industrial workcell environments and manipulation tasks
The lab models real-world industrial automation tasks across specialized manufacturing workcells:
1. **Industrial Workcell 1: Automotive Battery Cell Assembly & QC**: Cleanroom assembly workcell containing prismatic battery cells, polyamide insulator bushings, and aluminum busbar terminals (`industrial_battery.xml`).
2. **Industrial Workcell 2: Semiconductor Wafer Carrier & Fixture Kitting**: Cleanroom workcell with a safety barrier partition where 3D LiDAR ToF scanning guides high-elevation obstacle clearance (`industrial_semiconductor.xml`).
3. **Industrial Workcell 3: Turbine Bracket Machining with Distractor Perturbation**: Aerospace manufacturing cell handling hot-forged Titanium brackets and tool steel couplings while random distractor chips shift under coolant spray (`industrial_machining.xml`).
Corresponding long-horizon industrial tasks evaluate hierarchical composition:
- **Task 1: Automotive Battery Module Tri-Part Assembly (Raman QC Guided)**:
- *Stage 1*: Excitation Raman scan verifies polyethylene dielectric separator integrity (peaks at $1440, 2850-2920\,\text{cm}^{-1}$); manipulator transfers cell to battery chassis pocket.
- *Stage 2*: Raman scan detects amide I and II bands ($1550, 1640\,\text{cm}^{-1}$) verifying high-temperature polyamide insulator; seats part into isolation bay.
- *Stage 3*: Raman scan confirms aluminum 6061 alloy lattice phonon modes; docks busbar terminal into the terminal clamping dock.
- **Task 2: Semiconductor Wafer Carrier & 3D LiDAR Obstacle Traversal**:
- Manipulator scans vertical cleanroom fixture barriers with 3D LiDAR, executing an elevated clearance trajectory over the barrier wall into the vacuum chuck nest, while monitoring 77 GHz Radar Doppler velocity feedback.
- **Task 3: Turbine Bracket Machining with Distractor Perturbation**:
- Tests selective temporal abstraction in heavy machining workcells. The planner transfers hot-forged titanium brackets into a precision CNC hydraulic vise while adjacent distractor parts shift randomly without corrupting the macro plan.
## Empirical findings
### Selective abstraction across hierarchy depth
A critical property of hierarchical world models is selective temporal abstraction: higher levels discard rapid, unpredictable detail while preserving slower, task-relevant state variables.
Linear probes trained on frozen latent representations demonstrate this separation:
- In Visual AntMaze, probe error for fast-varying ant leg joint angles and distractor objects rises sharply at Level 2 and Level 3.
- In contrast, probe error for the global center-of-mass coordinates and maze position remains near zero across all hierarchy levels.
<div class="hjepa-video-grid">
<div class="hjepa-video-card">
<video controls autoplay muted loop playsinline preload="metadata">
<source src="media/hjepa/videos/selective_abstraction_lvl4.mp4" type="video/mp4">
</video>
<div class="hjepa-video-caption">
<strong>Selective Abstraction in AntMaze:</strong> Level 4 latent representations retain the spatial navigation manifold while discarding high-frequency leg oscillation artifacts.
</div>
</div>
<div class="hjepa-video-card">
<video controls autoplay muted loop playsinline preload="metadata">
<source src="media/hjepa/videos/scene-diversity.mp4" type="video/mp4">
</video>
<div class="hjepa-video-caption">
<strong>Real-World Scene Diversity:</strong> Comparison of fixed laboratory setups against DROID trajectories with varied objects, backgrounds, and camera angles.
</div>
</div>
</div>
### Goal-distance surfaces in latent space
The learned latent geometry fundamentally alters the planning cost landscape. In single-timescale representations, Euclidean distance in latent space is corrupted by extraneous visual variations. In hierarchical world models, the higher-level latent space produces smooth, monotonic goal-distance basins oriented toward task progress:
### Planning Pareto frontier
Across FourRoom navigation, Visual AntMaze, Push-T manipulation, and OGBench Cube manipulation, hierarchical world models achieve superior success rates per unit of planner compute compared to flat world models (LeWM):
On Visual AntMaze, adding a three-level hierarchy raises task completion from 18% to 73%.
<div class="hjepa-video-grid">
<div class="hjepa-video-card">
<video controls autoplay muted loop playsinline preload="metadata">
<source src="media/hjepa/videos/ant_ep29.mp4" type="video/mp4">
</video>
<div class="hjepa-video-caption">
<strong>Visual AntMaze Execution 1:</strong> Three-level hierarchical world model navigating complex maze junctions.
</div>
</div>
<div class="hjepa-video-card">
<video controls autoplay muted loop playsinline preload="metadata">
<source src="media/hjepa/videos/ant_ep35.mp4" type="video/mp4">
</video>
<div class="hjepa-video-caption">
<strong>Visual AntMaze Execution 2:</strong> Subgoal refinement avoiding deadlock in dead-end corridors.
</div>
</div>
</div>
<div class="hjepa-video-grid">
<div class="hjepa-video-card">
<video controls autoplay muted loop playsinline preload="metadata">
<source src="media/hjepa/videos/pusht_ep7.mp4" type="video/mp4">
</video>
<div class="hjepa-video-caption">
<strong>Push-T Rollout 1:</strong> Continuous planar precision alignment with coarse-to-fine trajectory planning.
</div>
</div>
<div class="hjepa-video-card">
<video controls autoplay muted loop playsinline preload="metadata">
<source src="media/hjepa/videos/cube_30_ep2.mp4" type="video/mp4">
</video>
<div class="hjepa-video-caption">
<strong>Cube Manipulation:</strong> Multimodal 3D contact-rich reorientation on OGBench.
</div>
</div>
</div>
### DROID real-robot trajectory fidelity
On diverse real-world manipulator demonstrations from DROID, inverse dynamics supervision ensures that the representation models end-effector trajectories rather than static background features.
Using Fréchet path fidelity against expert demonstrations, the hierarchical world model consistently outperforms flat baselines across varying computation budgets.
## Code implementation
The Python companion script [`examples/l23_9_hjepa.py`](code/examples/l23_9_hjepa.py) provides a complete, self-contained PyTorch implementation of a two-level hierarchical world model architecture with SIGReg regularization and a top-down latent visual planner:
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
class SIGRegLoss(nn.Module):
def __init__(self, lambd_mean=1.0, lambd_cov=1.0):
super().__init__()
self.lambd_mean = lambd_mean
self.lambd_cov = lambd_cov
def forward(self, z):
b, d = z.shape
mean = z.mean(dim=0)
mean_loss = torch.sum(mean ** 2)
z_centered = z - mean
cov = (z_centered.T @ z_centered) / (b - 1)
eye = torch.eye(d, device=z.device)
cov_loss = torch.sum((cov - eye) ** 2)
return self.lambd_mean * mean_loss + self.lambd_cov * cov_loss
class MultimodalSensorEncoder(nn.Module):
def __init__(self, token_dim=16):
super().__init__()
# Encoders for heterogeneous physical sensors: RGB-D, LiDAR, Radar, Raman, Proprio
self.rgb_enc = nn.Sequential(nn.Linear(32, 32), nn.SiLU(), nn.Linear(32, token_dim))
self.depth_enc = nn.Sequential(nn.Linear(16, 32), nn.SiLU(), nn.Linear(32, token_dim))
self.lidar_enc = nn.Sequential(nn.Linear(16, 32), nn.SiLU(), nn.Linear(32, token_dim))
self.radar_enc = nn.Sequential(nn.Linear(12, 32), nn.SiLU(), nn.Linear(32, token_dim))
self.raman_enc = nn.Sequential(nn.Linear(24, 32), nn.SiLU(), nn.Linear(32, token_dim))
self.proprio_enc = nn.Sequential(nn.Linear(14, 32), nn.SiLU(), nn.Linear(32, token_dim))
self.fusion = nn.Sequential(nn.Linear(6 * token_dim, 64), nn.SiLU(), nn.Linear(64, 32))
def forward(self, rgb, depth, lidar, radar, raman, proprio):
tokens = [self.rgb_enc(rgb), self.depth_enc(depth), self.lidar_enc(lidar),
self.radar_enc(radar), self.raman_enc(raman), self.proprio_enc(proprio)]
return self.fusion(torch.cat(tokens, dim=-1))
class HierarchicalJEPA(nn.Module):
def __init__(self, token_dim=16, act_dim=7, l1_dim=32, l2_dim=16, stride=4):
super().__init__()
self.stride = stride
self.sensor_encoder = MultimodalSensorEncoder(token_dim=token_dim)
self.pred1 = nn.Sequential(nn.Linear(l1_dim + act_dim, 64), nn.SiLU(), nn.Linear(64, l1_dim))
self.idm1 = nn.Sequential(nn.Linear(2 * l1_dim, 64), nn.SiLU(), nn.Linear(64, act_dim))
self.enc2 = nn.Sequential(nn.Linear(l1_dim, 32), nn.SiLU(), nn.Linear(32, l2_dim))
self.act_enc2 = nn.Sequential(nn.Linear(stride * act_dim, 32), nn.SiLU(), nn.Linear(32, act_dim))
self.pred2 = nn.Sequential(nn.Linear(l2_dim + act_dim, 32), nn.SiLU(), nn.Linear(32, l2_dim))
self.sigreg = SIGRegLoss()
def forward_loss(self, sensors, act_seq):
z1 = self.sensor_encoder(sensors["rgb"], sensors["depth"], sensors["lidar"],
sensors["radar"], sensors["raman"], sensors["proprio"])
z1_pred = self.pred1(torch.cat([z1[:, :-1], act_seq], dim=-1))
loss_pred1 = F.mse_loss(z1_pred, z1[:, 1:])
loss_sigreg1 = self.sigreg(z1.reshape(-1, z1.shape[-1]))
# Subsample coarse steps according to temporal stride
z2 = self.enc2(z1[:, ::self.stride])
macro_act = self.act_enc2(act_seq[:, :self.stride].reshape(act_seq.shape[0], -1))
z2_pred = self.pred2(torch.cat([z2[:, 0], macro_act], dim=-1))
loss_pred2 = F.mse_loss(z2_pred, z2[:, 1])
loss_sigreg2 = self.sigreg(z2.reshape(-1, z2.shape[-1]))
return loss_pred1 + 0.1 * loss_sigreg1 + loss_pred2 + 0.1 * loss_sigreg2
{
"questions": [
{
"prompt": "Why does training a single-timescale action-conditioned JEPA fail on long-horizon navigation and manipulation tasks?",
"options": [
"A single representation cannot simultaneously resolve high-frequency local physical interactions and slow task progress without accumulating predictive drift.",
"Single-timescale JEPAs cannot utilize Transformer attention heads.",
"JEPA architectures require pixel reconstruction loss to evaluate multi-step dynamics.",
"Contrastive negative samples become mathematically impossible to compute beyond ten timesteps."
],
"answer": 0,
"explanation": "When forced into a shared single latent space, rapid details (such as leg joint oscillations or local contact flutter) interfere with measuring global task progress, causing long-horizon prediction error to compound."
},
{
"prompt": "What role does SIGReg (Skew-Independent Gaussian Regularization) serve in the hierarchical world model objective?",
"options": [
"It acts as an inverse kinematics solver for the robot end-effector.",
"It forces latent embeddings to match an uncorrelated standard normal distribution, preventing dimensional representation collapse without negative pairs.",
"It synthesizes synthetic RGB video frames to supervise the observation decoder.",
"It penalizes high motor torque commands to guarantee smooth robotic motions."
],
"answer": 1,
"explanation": "SIGReg penalizes deviation from zero mean and identity covariance, ensuring that the latent coordinates do not collapse to trivial constants or redundant subspaces."
},
{
"prompt": "Why is an Inverse Dynamics Model (IDM) loss necessary when pretraining hierarchical world models on diverse real-world robot datasets such as DROID?",
"options": [
"To replace the forward predictive Transformer model entirely.",
"To prevent the latent representation from encoding the static background while ignoring the moving manipulator.",
"To calculate joint friction coefficients inside the physical hardware controllers.",
"To enforce that all demonstration trajectories have identical video durations."
],
"answer": 1,
"explanation": "In diverse robotic scenes, background pixels often dominate the image. Without an inverse dynamics loss predicting the action between consecutive frames, the encoder can achieve low prediction error by modeling the invariant background while discarding the manipulator."
}
]
}