Hierarchical, Hybrid, and Omnimodal Models

Michael BrenndoerferJune 30, 202658 min read

Part of World Models Handbook

Examines hierarchical, hybrid, and omnimodal world models, covering temporal abstraction, neural-symbolic hybrids, multimodal fusion.

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Hierarchical, Hybrid, and Omnimodal Models

A household robot wakes at 7:00 AM, unloads the dishwasher, hears you ask for coffee, and gets to work. It plans a sequence of steps that unfold over minutes: find the mug, walk to the coffee maker, fill the reservoir, grind beans, brew. Each step is executed by a fast motor-control loop that issues torque commands on a millisecond timescale. Along the way it fuses sight, sound, touch, and proprioception into a coherent picture of the kitchen. Different subsystems handle navigation, manipulation, and language, passing control as the task changes.

Consider how much that scenario asks of the complete agent. Its planner or task memory must preserve a goal for minutes while its controllers react to millisecond-scale events. It typically benefits from object-level or relational information for planning and continuous contact information for contact-rich manipulation, whether those signals are explicit or encoded across cooperating components. A termination rule or policy must detect when a subtask is finished and decide when to hand control between high-level skills and lower-level controllers. The system must also detect and respond safely when a camera is occluded or a microphone becomes uninformative. Trained missing-modality inference, filtering, and validated fallback behavior can improve tolerance to missing inputs, but do not by themselves guarantee safety. A monolithic fixed-cadence architecture can be inefficient and difficult to train for this combination, which motivates explicit abstraction and modularity.

Single-level world models, such as one recurrent state-space model from Part V, Ch 1 or one autoregressive predictor, often face difficult efficiency and error-propagation tradeoffs on this scenario when they lack explicit temporal abstraction. Temporal depth, representation type, modality fusion, and parameter sharing are distinct but interacting design axes: a temporally flat model can use discrete tokens, modality-specific encoders, or routed experts, while a hierarchy can still use continuous states and shared parameters. The architectural families in this chapter address those axes separately, though choices along one axis can constrain compute, interfaces, optimization, and useful choices along the others. A fixed-rate model may repeatedly recompute expensive slow context. An unstructured continuous latent can encode a fact such as "the mug is on the table," but does not expose relations and logical constraints as directly as a symbolic representation. Naive early fusion asks one encoder to accommodate signals with different scales, rates, and noise. Fully shared parameters can also invite interference across tasks, though flat models may instead use task heads, adapters, or experts.

Long-horizon manipulation, sim-to-real robotics, multimodal perception, and multi-task control expose overlapping limitations that can motivate temporal abstraction, hybrid modeling, modality-aware fusion, or modular parameters. None of these settings maps exclusively to one response. The useful point is that the responses compose: they are reusable ingredients rather than four mutually exclusive designs.

This chapter develops four architectural families that together lift those limitations:

  • Hierarchical models organize states or computation across temporal, spatial, or abstraction levels; a given design need not use every kind of scale.
  • Hybrid models combine learned components with symbolic reasoning or with hand-crafted simulators, using each where it is strongest.
  • Omnimodal models in this chapter support broad or any-to-any combinations of supported input and output modalities; the term is not yet standardized across the literature. Shared latent states and graceful missing-modality behavior are useful designs, not consequences of the label alone.
  • Modular models divide computation among components. Mixture-of-experts models are one modular family, using learned routers; other modular systems use fixed or symbolic composition.

These four families are not rivals. Modern systems can stack them: a hierarchical planner whose slow level uses a symbolic simulator and whose fast level uses sensor-specific experts around a multimodal latent. The rest of this chapter builds intuition for each, then compares a window-conditioned context model with a separate routed model.

As you read, it helps to keep three questions in mind for every architecture. What assumption of the flat model does it relax? What new failure mode does it introduce in exchange? And what quantity should we measure to tell whether it is helping? Most of the practical wisdom in this area lives in those three questions rather than in the architectures themselves.

Multiple Temporal and Spatial Scales

A useful household-robot architecture must span multiple functional timescales and often benefits from predictions and decisions at several levels of abstraction, whether those levels are explicit, implicit, or distributed across cooperating models and controllers. "What will the room look like in thirty seconds if I keep pouring?" is a slow, coarse prediction. "What torque should the wrist apply this millisecond?" is a fast control decision. The relevant state changes at different rates. A fixed-rate model can waste computation if it repeatedly recomputes expensive slow context at the fast control rate, and recursive prediction exposes it to compounding rollout error.

Recall from Part I, Ch 6 that when we roll a learned model forward, each predicted state becomes the input to the next prediction. Consider true and predicted rollouts that start from the same state and use the same action sequence. Suppose the transition is uniformly LL-Lipschitz in state over the rollout region and the one-step model discrepancy is uniformly bounded by ε\varepsilon in the same norm. Their error then obeys eh+1≤Leh+εe_{h+1} \leq L e_h + \varepsilon with e0=0e_0 = 0. Under these conditions, the worst-case envelope can remain bounded, grow linearly, or grow geometrically across HH steps. If the transition amplifies differences, the bound compounds geometrically; at unit gain it grows linearly; if the map is contractive, it saturates. A fast fixed-step model chains more learned transitions over a fixed physical horizon, though cadence is not itself the cause of the error. With initial error e0>0e_0 > 0, the same bound also contains the propagated term LHe0L^H e_0.

The compounding effect is easiest to see directly. The next cell solves that recurrence for a few values of LL, holding the uniform per-step discrepancy fixed and using e0=0e_0 = 0.

In[3]:
Code
import numpy as np

compound_horizon = np.arange(0, 31)
compound_epsilon = 0.01
compound_lipschitz = np.array([0.85, 1.0, 1.10])
compound_envelopes = []
for L in compound_lipschitz:
    if abs(L - 1.0) < 1e-12:
        vals = compound_epsilon * compound_horizon
    else:
        vals = compound_epsilon * (L**compound_horizon - 1.0) / (L - 1.0)
    compound_envelopes.append(vals)
compound_envelopes = np.stack(compound_envelopes, axis=0)
Out[4]:
Visualization
Three bounds for rollouts that start from the same state and use the same actions under uniform one-step discrepancy: the L=0.85 curve saturates, L=1 grows linearly, and L=1.10 bends upward geometrically.
Worst-case rollout error bounds under the stated same-state, same-action, uniform-discrepancy, and uniform-Lipschitz assumptions. At unit gain, errors accumulate linearly; above unit gain, they are amplified geometrically; below unit gain, the bound saturates.

Temporal hierarchy reduces the frequency of high-level goal or option decisions while retaining fast low-level control. If a subgoal latent updates every KK fast steps, a five-minute plan at one hundred control updates per second can contain three hundred one-second macro-decisions instead of thirty thousand primitive decisions. This shortens the high-level search chain. It can reduce compounding prediction error when coarse-transition error remains small enough and fine-level error is controlled, for example by local rollouts or regular correction from observations; hierarchy does not guarantee an empirical advantage.

The Semi-Markov View

A canonical reinforcement-learning formalization of temporal abstraction is the options framework of Sutton, Precup, and Singh (1999). The idea is to let the agent choose not just individual actions but whole temporally extended behaviors, and to reason about those behaviors as if they were single actions. Let S\mathcal{S} and A\mathcal{A} denote the underlying MDP's state and action spaces. An option is a triple (I,π,β)(\mathcal{I}, \pi, \beta) where:

  • I⊆S\mathcal{I} \subseteq \mathcal{S} is the set of states in which the option may be started.
  • π(a∣s)\pi(a \mid s) is the option's internal action distribution: π(a∣s)∈[0,1]\pi(a \mid s) \in [0,1] and ∑a∈Aπ(a∣s)=1\sum_{a \in \mathcal{A}} \pi(a \mid s)=1 for each ss.
  • β:S→[0,1]\beta: \mathcal{S} \to [0,1] is the termination probability.

Read the triple as a small policy with a guard rail. The initiation set I\mathcal{I} says where the option makes sense, the internal policy π\pi says what to do while it runs, and the termination probability β\beta says when to stop. When the option is chosen at time tt, it executes its internal policy until β\beta fires. Assuming the option terminates almost surely, its duration is a positive integer-valued random variable τ∈{1,2,… }\tau \in \{1, 2, \dots\}; depending on the option and environment, τ\tau may be stochastic or degenerate at a fixed duration. This turns the Markov decision process (MDP) into a semi-Markov decision process (SMDP), where the value of taking option oo in state ss is

Q(s,o)=E ⁣[∑i=t+1t+τγi−t−1ri  +  γτV(st+τ) | st=s,ot=o]Q(s, o) = \mathbb{E}\!\left[\sum_{i=t+1}^{t+\tau} \gamma^{i-t-1} r_i \;+\; \gamma^{\tau} V(s_{t+\tau}) \,\middle|\, s_t = s, o_t = o \right]

where:

  • Q(s,o)Q(s,o): the expected return from taking option oo in state ss.
  • ss: the current state.
  • sts_t: the state at time step tt.
  • oo: the option.
  • oto_t: the option chosen at time step tt.
  • τ\tau: the random duration of the option.
  • tt: the current time step.
  • γ\gamma: discount factor.
  • rir_i: reward at time step ii.
  • ii: time index running from t+1t+1 to t+τt+\tau.
  • VV: continuation value under the policy followed after the option terminates.
  • st+τs_{t+\tau}: the state after option oo terminates.

Here the expectation is over the option's internal action choices, the environment's transitions and rewards, and the random termination time.

Compare this to an ordinary one-step action value, which contains one immediate reward plus one continuation value discounted by γ\gamma. The option-value function instead sums the rewards earned during the whole option and then discounts the continuation value by γτ\gamma^{\tau}, the discount over the option's realized duration. Different options therefore carry different effective continuation discounts, and the high-level policy chooses temporally extended actions rather than atomic ones. For 0<γ<10<\gamma<1, continuation value after a longer realized duration receives a smaller weight.

The effective discount for an option is γτ\gamma^{\tau}, so duration directly changes the weights the high-level policy compares.

In[5]:
Code
import numpy as np

option_durations = np.arange(1, 21)
discount_rates = np.array([0.90, 0.95, 0.99])
option_discounts = np.stack(
    [gamma**option_durations for gamma in discount_rates], axis=0
)
Out[6]:
Visualization
Three monotonically decreasing curves of gamma raised to duration; smaller gamma values decline much faster.
Effective discount applied to an option's continuation value as a function of its realized duration. Longer options are discounted more heavily, so the high-level policy compares options at different effective horizons.

Getting the discounting right is essential because an incorrect duration discount can distort option values whenever it changes the weight on a nonzero continuation or reward term. For example, replacing γτ\gamma^\tau with γ\gamma overweights positive continuation value after long options; the direction of the resulting preference also depends on within-option rewards and the sign of the continuation value.

Temporal Abstraction

Temporal abstraction treats sequences of primitive actions as single high-level actions. It is a common way to extend the effective high-level planning horizon because one abstract decision can cover many primitive steps.

Two-Level and Multi-Level Latent Architectures

Several hierarchical world-model architectures implement temporal abstraction with multiscale latent states, latent skills, or latent goals rather than literal call-and-return option triples. One two-level generative design stacks latent variables that update at different rates:

  • A slow latent zksz^{s}_k that changes once every KK fast steps, indexed by macro-time kk.
  • A fast latent ztfz^{f}_t that updates every step and receives z⌊t/K⌋sz^{s}_{\lfloor t/K \rfloor} as conditioning context.

When training makes it encode persistent behavior-level context, the slow latent can play a role analogous to option identity or high-level intent. The fast latent updates each step and can represent variation that the persistent context does not explain. Conditioning the fast transition on the current slow latent lets its predictions specialize to that context.

The joint generative factorization is

p(z1:Tf,z1:Ls∣z0f,z0s,a0:T−1)=∏k=0L−1[p(zk+1s∣zks)∏t=kK(k+1)K−1p(zt+1f∣ztf,at,z⌊t/K⌋s)]p(z^{f}_{1:T}, z^{s}_{1:L} \mid z^{f}_0, z^{s}_0, a_{0:T-1}) = \prod_{k=0}^{L-1} \left[ p(z^{s}_{k+1} \mid z^{s}_k) \prod_{t = kK}^{(k+1)K - 1} p(z^{f}_{t+1} \mid z^{f}_t, a_t, z^{s}_{\lfloor t/K \rfloor}) \right]

where:

  • z1:Tfz^{f}_{1:T}: sequence of fast latent states over fast time steps t=1,…,Tt = 1, \dots, T.
  • z1:Lsz^{s}_{1:L}: sequence of slow latent states over slow time steps k=1,…,Lk = 1, \dots, L.
  • z0fz^{f}_0: initial fast latent state before the sequence.
  • z0sz^{s}_0: initial slow latent state.
  • LL: number of slow steps (macro-steps).
  • KK: number of fast steps per slow step, so T=L×KT = L \times K.
  • ata_t: action at fast time step tt; the joint distribution is conditioned on the action sequence a0:T−1a_{0:T-1}.
  • pp: probability distribution.
  • zksz^{s}_k: the slow latent state at macro-time step kk.
  • ztfz^{f}_t: the fast latent state at fast time step tt.
  • ⌊⋅⌋\lfloor \cdot \rfloor: the floor function, giving the macro-time index for fast time tt.

The outer product over kk runs over macro-steps, and the inner product over tt runs over the KK fast steps inside each macro-step. Reading the factorization from left to right, the slow equation defines the temporal coarsening; the fast equation defines the fine-grained dynamics conditioned on the current slow context. Notice how the slow transition p(zk+1s∣zks)p(z^{s}_{k+1} \mid z^{s}_k) does not depend on the actions at all in this minimal form: it is a coarse drift model of how the sub-goal evolves, while the fast level absorbs the effect of actions. In richer variants the slow transition is also conditioned on an abstract action or goal, but the essential split is that the slow level summarizes and the fast level reacts.

Encoders can mirror this structure. A training-time smoothing encoder q(zks∣okK:(k+1)K−1)q(z^{s}_k \mid o_{kK : (k+1)K - 1}) aggregates an entire observation window, while a fast encoder q(ztf∣ot,z⌊t/K⌋s)q(z^{f}_t \mid o_t, z^{s}_{\lfloor t/K \rfloor}) captures per-step detail. Here oto_t is the observation at fast time step tt, and okK:(k+1)K−1o_{kK:(k+1)K-1} is the corresponding window for slow step kk. Because this smoothing form uses the completed window, it is not a causal online encoder for early steps in that same block. An online system can instead infer zksz^s_k from the preceding completed window o(k−1)K:kK−1o_{(k-1)K:kK-1} before block kk, or delay use of the summary until its window is complete. In either case, the asymmetry is the point: the slow encoder compresses a batch of observations, while the fast encoder sees a single observation together with the summary currently in force. A useful analogy is human scene perception: an initial glance can provide scene gist that helps guide subsequent fixations toward task-relevant detail.

FeUdal Networks (Vezhnevets et al., 2017) use a lower-temporal-resolution Manager whose dilated recurrent state and multi-step objective produce persistent goal conditioning, while a Worker chooses primitive actions each step. The Manager and Worker are learned concurrently with separate objectives, without propagating gradients between the two modules. Director adds hierarchical policies over the latent state space of a learned pixel-based world model: a Manager selects latent goals every eight steps and a Worker chooses primitive actions to reach them. These are policy hierarchies rather than evidence that Director contains slow and fast world-model state transitions. More generally, persistent context can spare a fast predictor from rediscovering behavior-level information at every step, but smaller residuals are an empirical outcome rather than a guarantee.

Spatial Scales

Temporal scales have a spatial counterpart. A building-scale robot may use room-level topology for navigation together with millimeter-scale contact geometry and high-rate force or torque signals for manipulation. With a fixed-budget dense representation, one resolution is inefficient across these scales: a fine global grid allocates detail where it is irrelevant, while a coarse grid misses local contact structure. A dense volumetric representation of an entire building at one-millimeter resolution is computationally impractical, so sparse or hierarchical representations allocate fine resolution selectively. Two common strategies are:

  • Coarse-to-fine latent pyramids. The state has a low-resolution "global" component, for example a volumetric or object-level map, plus fine local patches that are only queried where the agent is currently interacting. The global component ignores details that do not matter; the local component ignores the rest of the world. The two are linked by the agent's current focus, so the fine patches are always embedded in a coarse frame of reference.
  • Multi-resolution observation models. The decoder reconstructs a coarse observation everywhere and a fine observation only inside an attention window around the agent. Biological foveation offers an analogy: cone density and visual acuity peak in a small central region, while gaze shifts sample the rest of the scene. Most retinal photoreceptors, predominantly rods, are outside the fovea.

Structure, not resolution alone, should drive the transition model. Regular-grid maps often suit convolution because of spatial locality and translation equivariance; topological maps often suit graph networks because they encode adjacency and relations. Continuous-time models can help with irregularly timed observations or queries, while contact dynamics may additionally require event-aware modeling and accurate numerical integration. Graph networks can themselves be effective for contact dynamics. The practical rule is to match an inductive bias to the representation's actual topology, smoothness, events, and sampling process.

Failure Modes

Hierarchical models add complexity, and that complexity has costs. A slow policy can propose subgoals that a fast policy cannot realize. Termination can be poorly calibrated, causing premature or excessively long execution. When observations support several contexts, slow-state inference can become uncertain or lock onto the wrong one unless history, actions, or later observations correct it. A window encoder can also exploit spurious phase or boundary cues. Window length itself becomes a shortcut when variable lengths or padding patterns correlate with the target. Such correlations fail when task periods or window construction change.

Several of these failure modes can manifest as disagreement between the levels, with no training signal forcing them to reconcile. Silent inconsistency is especially difficult because it may not raise an obvious error. The slow level believes the robot is near the mug, the fast level believes it is near the sink, and the model produces plausible but wrong predictions for both. This drift can remain hidden from aggregate one-step validation loss, so cross-level agreement should be evaluated explicitly or included as an equivalent targeted validation term.

Cross-level consistency losses address the disagreement subset, not every failure above. If the slow level predicts a projected coarse summary of the next window and the fast level predicts its fine state sequence, applying the same projection to the fast prediction gives a shared quantity on which the two levels can be compared. Training on that agreement supplies a direct gradient to the slow level. Subgoal feasibility instead calls for reachability or capability checks; termination needs dedicated closed-loop stopping metrics; and window shortcuts need distribution-shift tests and ablations. Joint end-to-end training and a bounded window matched to the task's timescale may support these safeguards, but do not replace them.

Neural-Symbolic and Simulator Hybrids

Pure learned models offer flexible function approximation, but their data requirements depend on architecture, priors, pretraining, and task complexity. With sufficient capacity, suitable data, and successful optimization, they can approximate a broad range of patterns, though they do not automatically enforce logical consistency. Symbolic systems can enforce their encoded rules exactly, while hand-crafted simulators can be accurate within their assumptions and numerical tolerances. Both become brittle when those assumptions fail. Hybrid models combine these strengths: learned components can handle perception, uncertainty, and difficult-to-model residuals, while symbolic or simulation components can implement explicit planning rules, logical constraints, and well-understood physics. The central design choice is where to draw the boundary and how information should cross it.

Three Integration Patterns

There is no single way to combine neural and symbolic components. In practice the integrations fall into three broad patterns, distinguished by how much information flows between the learned and non-learned parts and in which direction.

Pipeline hybrids. One common pipeline hybrid uses learned perception to map raw observations into a symbolic state of typed objects, predicates, and relations. In the neural-symbolic pipeline described here, a symbolic reasoner plans over that state and a learned controller executes the plan; other pipeline hybrids may use classical or hand-designed controllers. Shakey was an early precursor to this perception-model-planning-execution pattern, and some modern task-and-motion planning systems use learned perception or predicates as a front end. Interfaces between modules are discrete and non-differentiable: a planning cycle may run from perception to planner to controller, while execution feedback triggers new perception and replanning. The symbolic layer makes assumptions inspectable. The cost is that downstream plan failure supplies no end-to-end gradient to correct a perception mistake.

Differentiable hybrids. The symbolic component is made differentiable (or a soft relaxation is used) so gradients can flow through it. Differentiable planners, differentiable logic layers, and differentiable physics simulators all live here. The advantage is end-to-end learning: because the gradient reaches the symbolic part, the learned perception can be shaped by the downstream planning objective, closing the loop that the pipeline pattern leaves open. The disadvantage is that soft relaxations can be approximate and the optimizer may drift into regimes where the relaxation is unfaithful. A differentiable planner that relaxes discrete decisions into continuous weights might, during training, settle into a continuous solution that has no valid discrete counterpart, and the resulting policy can be brittle at deployment.

Residual hybrids. A hand-crafted model supplies a nominal prediction, and a neural network predicts an additive correction; when the nominal model is sufficiently accurate, it may account for most of the final prediction. During training, the residual is fit to the model's mismatch. This pattern sidesteps the optimization hazards of the differentiable hybrid by treating the non-learned part as fixed: the simulator is evaluated in the forward direction and the network learns its prediction error. Multi-step use is open loop only when predicted states are recursively fed back without observation correction. Let fsimf_{\text{sim}} be the simulator and gθg_\theta the neural correction. The hybrid dynamics are

s^t+1=fsim(st,at)+gθ(st,at),\hat{s}_{t+1} = f_{\text{sim}}(s_t, a_t) + g_\theta(s_t, a_t),

and the training loss is

L(θ)=E(st,at,st+1)∥st+1−fsim(st,at)−gθ(st,at)∥22.\mathcal{L}(\theta) = \mathbb{E}_{(s_t, a_t, s_{t+1})} \big\| s_{t+1} - f_{\text{sim}}(s_t, a_t) - g_\theta(s_t, a_t) \big\|_2^2.

where:

  • s^t+1\hat{s}_{t+1}: the predicted next state.
  • fsim(st,at)f_{\text{sim}}(s_t, a_t): the simulator's predicted next state.
  • gθ(st,at)g_\theta(s_t, a_t): the neural residual correction.
  • θ\theta: parameters of the residual correction network.
  • tt: time step index.
  • sts_t: the current state.
  • ata_t: the action at time tt.
  • st+1s_{t+1}: the true next state.
  • L(θ)\mathcal{L}(\theta): the mean squared residual loss used to train the correction network.
  • E(st,at,st+1)\mathbb{E}_{(s_t, a_t, s_{t+1})}: expectation over the training distribution of transitions.
  • ∥⋅∥22\|\cdot\|_2^2: squared Euclidean norm.

Read the loss as asking the network to explain only the part of the transition that the simulator misses. When the simulator's expected error is small relative to the scale of the chosen full-model target under the same normalization and distribution, the residual learner receives a smaller-magnitude target. Residual learning is generally more favorable when the remaining mismatch is predictable, observable, and representable by the learner.

The practical benefit is that the residual network can see a smaller target. The next cell isolates that scale effect by simulating residual targets for an accurate simulator and a biased simulator while holding the true transition distribution fixed. It does not test whether either residual is predictable from the model inputs. The sample is drawn with seed 20240517.

In[7]:
Code
import numpy as np

residual_rng = np.random.default_rng(20240517)
n_residual_samples = 4000
true_transitions = residual_rng.normal(size=(n_residual_samples, 4))
# Accurate simulator: small-noise residual
true_transition_shift = 0.0 * residual_rng.normal(size=true_transitions.shape)
accurate_simulator = (
    true_transitions
    + true_transition_shift
    + 0.15 * residual_rng.normal(size=true_transitions.shape)
)
# Biased simulator: systematic offset increases residual magnitude
biased_simulator = (
    true_transitions
    + 0.35
    + 0.80 * residual_rng.normal(size=true_transitions.shape)
)
accurate_residual_norms = np.linalg.norm(
    true_transitions - accurate_simulator, axis=1
)
biased_residual_norms = np.linalg.norm(
    true_transitions - biased_simulator, axis=1
)
Out[8]:
Visualization
Two histograms: accurate-simulator residual norms cluster near zero, while biased-simulator residual norms extend much farther right.
Distribution of residual target norms for an accurate simulator and a biased simulator. The accurate-simulator distribution is concentrated near zero, while the biased simulator leaves a broader, larger correction target.

Residual hybrids are used in robotics when a calibrated simulator captures much of the target dynamics and a learner can model the remaining friction, contact, or deformation effects. Under those conditions, gθg_\theta receives a smaller target than a model trained to predict the full transition from scratch. The gain depends on the base simulator and on whether its mismatch is observable and representable.

Differentiable Simulators

A differentiable simulator exposes gradients of simulated states with respect to control inputs and physical parameters. Instead of training a separate learned model, the simulator can serve as the world model for well-understood parts of the scene. Examples include DiffTaichi, Dojo, legacy Brax physics pipelines, and MuJoCo's current MJX-JAX implementation; MJX-Warp does not currently support automatic differentiation. These gradients support parameter fitting and, when reliable, end-to-end policy or planner optimization through simulation.

The cost is that simulator gradients can be biased or numerically delicate near contacts and other discontinuities. An infinitesimal input change may alter the contact sequence, producing a gradient that is uninformative for the intended update. When contact gradients are unreliable, a residual or other forward-only hybrid that treats the simulator as fixed avoids backpropagating through those gradients and may be preferable; its reliability still has to be established for the particular simulator and task.

Symbolic State and Learned Predicates

The relational planning representation used here describes entities through predicates such as on(mug, table), inside(bean, grinder), and reachable(gripper, handle); a domain may also assign entity types. Neuro-symbolic systems can learn such predicates from data. A perception encoder embeds detected objects, and a relation head produces soft truth values that a planner can query. A differentiable training path can use those soft values, while a discrete planner may threshold them at inference time.

The key modeling benefit is systematic generalization. Because lifted rules are stated over variables, a symbolic planner can apply them to unseen object combinations when the symbolic state is accurate, all predicates, types, and constraints hold, its operators cover the case, and search succeeds. Many standard neural architectures perform poorly on controlled compositional splits unless their architecture, objective, or training distribution supports systematic recombination. As covered in Part IV, Ch 6, explicit relational structure is one way to encode that support.

Neural-Symbolic Hybrid

In the coupled neural-symbolic world-model design used here, the state contains continuous latents for perception and control together with discrete or soft symbolic facts for planning and constraint checking. The two representations may be trained jointly, and either side may condition on the other; pipeline systems may instead train their modules separately.

Simulator-in-the-Loop World Models

A practical hybrid pattern is simulator-in-the-loop training. When real interaction is expensive, simulation can supply much of the training data while limited real data supports calibration, residual learning, or validation. A learned residual can reduce systematic simulator error; domain randomization instead trains the policy to tolerate a family of simulator variations. The physics engine supplies encoded structure, and the learned component addresses selected gaps. A corrected model may reduce exploitable bias, but policy optimization can still seek its remaining errors. Uncertainty penalties, robust optimization, constrained rollouts, and real-system validation are separate safeguards against that failure.

Shared State Across Modalities

A robot may receive vision, depth, touch, proprioception, audio, and language. A shared-latent multimodal world model gives those signals a common predictive state, allowing one modality to help disambiguate another. Separate private streams can also exchange information through cross-attention, so one literal shared bottleneck is not the only solution. Tolerance to missing inputs likewise requires suitable architecture and training; the label omnimodal does not guarantee it. The analogy to classical filtering is useful but limited: a Kalman filter combines measurements under an explicit probabilistic model, whereas a neural fusion rule and its state are learned.

Fusion Strategies

Three common fusion strategies are early fusion, late fusion, and intermediate exchange through mechanisms such as cross-modal attention. The first two describe placement near the input or output, while attention is an interaction operator that can be inserted at different depths and can coexist with either layout. Both placement and mechanism affect robustness and the kinds of correlations the model can capture.

Early fusion. A common design concatenates modality streams near the input and passes them through a shared encoder. Without explicit reliability weights, that encoder must infer each modality's influence during training. Masks, imputation, or fixed missing-value slots can make the design tolerate absent inputs; untrained omissions remain a distribution shift.

Late fusion. Each modality is encoded independently, and modality-specific predictions are combined at the output. Dropping an encoder is structurally simple, though the combiner still needs training for missing inputs. Because interaction occurs only after separate interpretations have formed, this design has less opportunity to model fine cross-modal correlations.

Intermediate exchange with cross-modal attention. In one common design, modalities receive separate preprocessing or encoders before attention exchanges information among their representations. Flamingo is a direct visual-to-language cross-attention example; Perceiver IO is a related multimodal architecture that routes arbitrary inputs through a latent bottleneck. This exchange can occur at one or several depths rather than defining a single placement category. Cross-modal attention adds interaction cost relative to independent encoders plus a simple combiner. Dense pairwise attention does not scale gracefully by itself, but sparse, asymmetric, or bottlenecked variants can. Attention masks make absent streams structurally representable, but reliable performance generally requires training on the relevant missingness patterns.

A shared latent ztz_t with MM modalities is inferred from

q(zt∣ot(1),…,ot(M))=q ⁣(zt  |  Fuse(e1(ot(1)),…,eM(ot(M)))),q(z_t \mid o_t^{(1)}, \dots, o_t^{(M)}) = q\!\left(z_t \;\middle|\; \text{Fuse}\big(e_1(o_t^{(1)}), \dots, e_M(o_t^{(M)})\big)\right),

where:

  • q(zt∣⋅)q(z_t \mid \cdot): the approximate posterior over the shared latent given the available modality observations.
  • ztz_t: shared latent state at time tt.
  • ot(m)o_t^{(m)}: observation from modality mm at time tt.
  • MM: number of modalities.
  • tt: time step index.
  • mm: modality index running from 11 to MM.
  • eme_m: per-modality encoder for modality mm.
  • Fuse\text{Fuse}: fusion function (concatenation, attention, or a learned pooling operator). The decoder mirrors this by producing modality-specific heads from the same ztz_t.

The equation says that each available modality is encoded and that a fusion operator maps the resulting embeddings into a shared latent. Modality-specific heads decode from ztz_t. The latent is a strict shared bottleneck only if it is the sole predictive path from inputs to outputs. Private or skip paths do not make the latent unshared, but they make the bottleneck nonexclusive.

Training Objectives

What forces the shared latent to be shared? Left to their own devices, per-modality decoders can send most information through private paths and never develop a common state. A pure reconstruction objective does not prevent this because each decoder only has to rebuild its own input. Four objectives encourage fusion:

  • Reconstruction. Each modality head must reconstruct its own input from ztz_t. Simple but weak on its own for encouraging cross-modal content, since it can be satisfied without any cross-modal interaction.
  • Cross-modal prediction. Predict one modality from the others via ztz_t: for example, recover depth from RGB and proprioception, or audio from video. This rewards content that is predictable across modalities. It cannot recover target information that is statistically absent from the inputs.
  • Contrastive alignment. Raise the similarity of matched cross-modal pairs relative to sampled negatives, thereby lowering the negatives' relative scores. A common formulation is the InfoNCE loss
LInfoNCE=−log⁡exp⁡(sim(zt(a),zt(b))/τ)exp⁡(sim(zt(a),zt(b))/τ)+∑j=1Nexp⁡(sim(zt(a),zj(b))/τ)\mathcal{L}_{\text{InfoNCE}} = -\log \frac{\exp(\text{sim}(z_t^{(a)}, z_t^{(b)})/\tau)}{\exp(\text{sim}(z_t^{(a)}, z_t^{(b)})/\tau) + \sum_{j=1}^{N} \exp(\text{sim}(z_t^{(a)}, z_j^{(b)})/\tau)}

where:

  • LInfoNCE\mathcal{L}_{\text{InfoNCE}}: the InfoNCE loss value, minimized when matching pairs have high similarity relative to negative samples.
  • zt(a)z_t^{(a)}: embedding of modality aa at time tt.
  • zt(b)z_t^{(b)}: embedding of modality bb at time tt.
  • zj(b)z_j^{(b)}: embedding of modality bb for the jj-th negative sample.
  • jj: negative-sample index.
  • NN: number of negative samples.
  • tt: time step index.
  • sim(⋅,⋅)\text{sim}(\cdot, \cdot): similarity function, usually cosine similarity.
  • τ\tau: temperature parameter (smaller τ\tau sharpens the distribution over negatives).

The numerator scores the positive pair, while the denominator sums the positive and negative scores, so minimizing the loss raises the positive pair's similarity relative to the negatives. The temperature τ\tau controls how sharply the loss penalizes near-miss negatives: a small τ\tau makes the model care about fine distinctions, while a large τ\tau makes it more forgiving.

  • Downstream compatibility. Train ztz_t as the state for a control or planning task. When the latent is the only predictive path, the task loss directly rewards retaining information needed for that task. It does not guarantee that every decision-relevant fact survives finite data, capacity, and optimization.

Practical systems often combine a subset of these objectives with task-specific weights. Reconstruction can provide a useful auxiliary signal, cross-prediction rewards predictable shared content, contrastive alignment organizes joint embeddings, and downstream tasks reward decision relevance. None guarantees a shared, stable, useful latent by itself, and the best combination depends on the modalities and downstream use.

Missing Modalities

An important empirical test is tolerance to missing inputs, though it measures ablation tolerance rather than proving that a latent is shared. At test time, inputs may be occluded, corrupted, or unavailable on a platform. Modality dropout exposes a model to many missingness patterns and can improve performance on covered or similar subsets. Architectures can also use fixed latent queries that attend to available tokens, or fixed modality slots with masks. Neither mechanism guarantees graceful degradation on unseen missingness patterns.

Empirical tolerance varies substantially with the dataset, task, fusion architecture, training modality set, and evaluation missingness pattern. In some trained systems, a task-specific subset of informative modalities preserves much of full-input performance, but the retained fraction depends on the dataset, task, architecture, and missingness training.

A simple reliability-weighted fusion model shows the shape of graceful degradation. The next cell computes the posterior standard deviation of a shared latent as modalities are removed in different orders.

In[9]:
Code
import numpy as np

modality_precisions = np.array([100.0, 50.0, 20.0, 5.0, 1.0, 0.5])
modality_available_counts = np.arange(1, len(modality_precisions) + 1)
modality_prior_precision = 0.01
modality_errors_keep_high = []
modality_errors_drop_high = []
for k in modality_available_counts:
    kept_prec_high = modality_precisions[:k]
    posterior_var_high = 1.0 / (modality_prior_precision + kept_prec_high.sum())
    modality_errors_keep_high.append(np.sqrt(posterior_var_high))
    kept_prec_low = modality_precisions[-k:]
    posterior_var_low = 1.0 / (modality_prior_precision + kept_prec_low.sum())
    modality_errors_drop_high.append(np.sqrt(posterior_var_low))
modality_errors_keep_high = np.array(modality_errors_keep_high)
modality_errors_drop_high = np.array(modality_errors_drop_high)
Out[10]:
Visualization
Two decreasing uncertainty curves: the curve for retained high-precision measurements remains below the low-precision curve until both meet when all six measurements are retained.
Posterior standard deviation in a scalar Gaussian reliability model as measurements are retained in different precision orders. For every retained count below six, keeping the assigned high-precision measurements yields lower posterior uncertainty than keeping the same number of low-precision measurements; the curves coincide when all measurements are retained.

This calculation assumes conditionally independent, unbiased Gaussian measurements of one scalar latent with known precisions. It shows that posterior uncertainty is dominated by the assigned high-precision measurements; it does not establish redundancy among real sensor modalities. Deployment sensor choices therefore require task-specific ablations rather than a general rule that fewer sensors will suffice.

Omnimodal World Models

In this chapter, an omnimodal world model means an any-to-any generative model designed to accept arbitrary supported input subsets and produce arbitrary supported output subsets. That definition is stronger than merely handling many modalities or learning a shared embedding. Gato and Unified-IO were early unified token-based multimodal systems: modality-specific serialization or tokenization feeds shared transformer processing. ImageBind instead aligns separate modality encoders in one embedding space and is not an arbitrary-modality generator. These systems anticipated parts of later any-to-any designs, but do not all satisfy the strong definition used here. In a dynamics configuration, a unified token model can predict trained future modality tokens from supported histories and actions.

Modular Experts and Routing

Increasing appropriately trained model capacity often improves performance, but larger dense models increase memory and inference cost and may be impractical to deploy. A sparse mixture-of-experts (MoE) layer can decouple total parameter count from per-example expert FLOPs by keeping many parameter sets but activating only a few per input. Parameter storage, routing, dispatch, communication, and memory traffic can still grow with the number and placement of experts. In world models the pattern is doubly useful: different regions of state space, different modalities, or different task contexts call for different dynamics, and a router can learn to select the right one. The first benefit is computational: a model that would be too slow to run in full can reduce expert computation by evaluating only a fraction of its experts per input. The second is an architectural inductive bias: specialization may reduce interference by assigning different regimes to different experts, although it does not automatically enlarge the representable function class relative to an unrestricted sufficiently large monolithic network.

The MoE Formulation

A mixture of KK experts with a learned router gg computes

p(y∣x)=∑k=1Kgk(x) fk(y∣x),gk(x)≥0,∑k=1Kgk(x)=1.p(y \mid x) = \sum_{k=1}^{K} g_k(x) \, f_k(y \mid x), \qquad g_k(x) \geq 0, \quad \sum_{k=1}^{K} g_k(x) = 1.

where:

  • p(y∣x)p(y \mid x): the output distribution over yy given input xx.
  • xx: input to the mixture.
  • yy: output variable.
  • KK: number of experts.
  • kk: expert index running from 11 to KK.
  • gk(x)g_k(x): router weight for expert kk given input xx.
  • fk(y∣x)f_k(y \mid x): the predictive density produced by expert kk for output yy given input xx.

The output density is a mixture of expert densities with input-dependent weights. Unlike an ordinary finite mixture with input-independent component proportions, an MoE uses a learned gating network whose weights depend on xx. Influential sparse Transformer MoEs commonly use a lightweight linear router, though more structured routers are possible:

g(x)=softmax(Wgx+bg).g(x) = \text{softmax}(W_g x + b_g).

where:

  • xx: input feature vector passed to the router.
  • g(x)g(x): router probability vector over experts.
  • WgW_g: weight matrix of the router.
  • bgb_g: bias vector of the router.
  • softmax\text{softmax}: softmax function.

The softmax turns raw scores into non-negative weights that sum to one, giving a valid mixture density. Keeping the router lightweight limits overhead. Specialization depends more directly on its inputs, objective, load balancing, capacity, and optimization than on router size alone.

Dense MoE computes all experts and combines them, so its expert compute grows with KK and it does not provide the conditional-compute scaling advantage of sparse MoE. A common sparse token-choice MoE selects the top kk experts (often k=1k = 1 or 22) and may renormalize weights over the selected set; other sparse routing schemes include top-1 Switch routing, expert-choice routing, threshold routing, and fixed or learned alternatives. Sparse MoE grows capacity (number of parameters) without growing per-example expert compute (which scales with kk, not KK). This decoupling is the central practical advantage: you can add experts to raise the model's total capacity while keeping the expert cost per example essentially flat, because only the selected experts run.

Routing Mechanisms

The router can be trained in several ways, and the choice trades off stability, compute, and gradient quality.

  • Learned linear gate. Simplest and most common. Backpropagates through g(x)g(x), so the router is trained jointly with the experts. Every expert receives a task gradient weighted by its router probability, giving continuous routing gradients, although saturated probabilities or optimization dynamics can still produce instability or negligible gradients.
  • Top-1 or top-k hard routing. The router still produces soft logits, but only the top-kk experts are evaluated and only those receive gradients. Sparse but potentially unstable, because hard assignments can starve experts of data: an expert that is never selected never receives a gradient and can never improve enough to be selected.
  • Noisy top-k. Add Gaussian noise to the router logits before top-kk selection so that different samples can explore different experts (Shazeer et al., 2017). With a nonzero noise scale large enough relative to the logit margins, this can reduce the risk that a small early advantage starves an expert of future updates. It may improve load distribution, but does not guarantee balanced traffic; explicit balancing objectives and capacity mechanisms are commonly used alongside it.
  • Soft relaxations. Dense soft routing computes all experts with continuous weights and trades expert compute for smoother routing gradients. Differentiable soft-top-k relaxations can instead be truncated to preserve sparse execution while approximating discrete selection.

Load Balancing and Expert Collapse

Without effective mitigation, learned sparse token-choice routers can develop severe load imbalance or collapse, with a few experts absorbing most traffic. The collapse can be self-reinforcing, because experts that start slightly ahead receive more data and improve faster, which makes them even more likely to be chosen next. A common mitigation is an auxiliary load-balancing loss. For the common top-1 Switch-style form, let nkn_k be the fraction of tokens routed to expert kk in a batch and let pkp_k be the average router probability for expert kk. Because each token receives one hard assignment, both vectors sum to one, and the auxiliary loss is

Lbalance=K∑k=1Knk⋅pk\mathcal{L}_{\text{balance}} = K \sum_{k=1}^{K} n_k \cdot p_k

where:

  • Lbalance\mathcal{L}_{\text{balance}}: auxiliary load-balancing loss used to discourage concentrated routing.
  • KK: number of experts.
  • kk: expert index running from 11 to KK.
  • nkn_k: fraction of tokens assigned to expert kk by top-1 routing in the batch.
  • pkp_k: average router probability for expert kk over the batch.

Uniform top-1 routing is a useful reference point. If nk=pk=1/Kn_k = p_k = 1/K for every expert, then

Lbalance=K∑k=1K1K1K=K(K1K2)=1\begin{aligned} \mathcal{L}_{\text{balance}} &= K \sum_{k=1}^{K} \frac{1}{K}\frac{1}{K} \\ &= K \left(K \frac{1}{K^2}\right) \\ &= 1 \end{aligned}

The loss is added to the task objective with a small coefficient. A concentrated router instead makes the hard token fractions and average soft probabilities large for the same experts, increasing their dot product.

The auxiliary loss has a simple shape as routing concentrates. The next cell computes the top-1 load-balancing loss for a four-expert router as the highest-weight expert receives an increasing fraction of the traffic.

In[11]:
Code
import numpy as np

balance_num_experts = 4
balance_focus_fractions = np.linspace(1.0 / balance_num_experts, 0.97, 60)
balance_loss_values = []
for focus_fraction in balance_focus_fractions:
    n = np.full(
        balance_num_experts, (1.0 - focus_fraction) / (balance_num_experts - 1)
    )
    n[0] = focus_fraction
    p = n.copy()
    balance_loss_values.append(balance_num_experts * np.sum(n * p))
balance_loss_values = np.array(balance_loss_values)
Out[12]:
Visualization
Line plot of load-balancing loss rising as routing concentration increases.
Top-1 auxiliary load-balancing loss as one expert absorbs a growing share of the routing traffic. Uniform top-1 routing gives the reference value one; concentration raises the loss because the hard and soft fractions agree on the same overused expert.

This auxiliary objective is a heuristic rather than a guarantee of uniform use because nkn_k and pkp_k are different statistics; implementations also rely on routing-capacity constraints or alternative objectives. If its coefficient is too small, concentrated routing may persist. If it is too large, the router can be pushed toward uniformity even when the data favors specialization.

Expert Collapse

The degenerate MoE regime in which a few experts absorb nearly all traffic while the rest remain undertrained. It wastes capacity and reduces effective parameter count. Sparse MoEs usually need a routing design, constraint, or auxiliary objective that prevents harmful collapse.

Modular World Models

Routing is not only about scaling. A decomposed world model can capture qualitatively different dynamics. Different regions of state space, different modes of a hybrid system, different tasks, or different modalities may benefit from distinct dynamics. A router that learns to send samples to the right expert effectively performs an implicit soft discretization of state space. Even when conditional-compute savings are not the main concern, explicit routing can provide a useful modular inductive bias and may improve optimization or sample efficiency by assigning different regimes to different experts. It does not automatically enlarge the representable function class relative to an unrestricted sufficiently large monolithic network.

This is related to the switched systems of classical control, where s˙=fσ(t)(s,a)\dot{s} = f_{\sigma(t)}(s, a) and σ(t)\sigma(t) identifies an active discrete mode. Here s˙\dot{s} is the state derivative, aa is the control, and fσ(t)f_{\sigma(t)} is mode-specific dynamics. A hard-routed MoE can be interpreted as learning a state/action-dependent mode selector; a soft MoE instead learns a weighted blend. Under a hard, mode-specialized interpretation, router chattering near a boundary is one relevant control concern.

Expert capacity should suit the complexity of the regime it handles. Near mode boundaries, routing and expert outputs should also avoid spurious jumps or rapid chattering while preserving real discontinuities in the system. Noise encourages routing exploration; entropy regularization discourages overly sharp probabilities; soft kernel-based routers can interpolate continuously. Output compatibility, hysteresis, or explicit continuity regularization may still be needed. The right treatment depends on whether the physical transition is continuous at that boundary.

Where Modularity Helps Most

  • Multi-embodiment robots. Different bodies can share perception while using embodiment-specific dynamics. If embodiment information is omitted or aliased in the represented input, a point-loss shared model can average incompatible body-specific transitions. Experts routed by embodiment are one remedy; conditioning a sufficiently expressive dense model on morphology is another.
  • Multi-task control. Navigation, manipulation, and speech respond to inputs differently; separate experts reduce interference. Task interference is the reason a single network trained on many tasks can lose accuracy on each of them, and routing sidesteps it.
  • Heterogeneous modalities. Different sensors have different statistics and noise; separate experts per modality concentrate compute where it is needed.
  • Regime-switching physical systems. Contact, free-fall, and sliding dynamics have qualitatively different structure. A learned soft boundary can track gradual transitions when the dynamics are continuous and the router is trained on the boundary region; discontinuous transitions may require hard, hysteretic, or event-aware treatment.

Worked Example: Comparing History Conditioning and Routing

We now compare a flat baseline with two separate window-conditioned alternatives on a controlled 2D system with a hidden context. One compresses a history window into a continuous context latent; the other uses the window to route across experts. The first borrows a representation idea from hierarchy but recomputes its context every step, so it is not a true multirate temporal hierarchy. The system has two contexts that change how actions affect velocity:

  • Context A is a damped point mass: each velocity component decays toward zero, and actions add acceleration.
  • Context B adds a rotational coupling term in velocity space, as if the agent is inside a rotating frame.

Context is fixed per episode and its label is not supplied to the models. A flat model sees only (state, action). A history-conditioned model and an MoE router each receive a short state window that can provide more statistical evidence about the context. Finite state history without past actions does not guarantee exact identification, and a single post-divergence state can itself carry probabilistic context information.

Both contexts agree at the initial rest state. After velocity becomes nonzero, subsequent transitions reveal their different update laws, and a longer window may add evidence. Two experts provide the cardinality needed for a possible one-to-one context alignment, but training must still discover and validate that alignment. Multi-step rollout error is therefore the main outcome, while routing statistics diagnose how the MoE uses its capacity.

Generating the Data

We build a NumPy rollout that produces an episode given the hidden context. The system has four state dimensions (position and velocity in 2D) and two action dimensions (accelerations).

In[13]:
Code
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(7)
rng = np.random.default_rng(7)

DT = 0.05
EPISODE_HORIZON = 40
WINDOW = 8
LATENT_DIM = 2
N_EXPERTS = 2
HIDDEN = 64
TRAIN_EPOCHS = 25
LEARNING_RATE = 1e-3
BATCH_SIZE = 256
EVAL_HORIZON = 20


def rollout(context, horizon=EPISODE_HORIZON, actions=None):
    if actions is None:
        actions = 0.5 * rng.standard_normal((horizon, 2)).astype(np.float32)
    else:
        actions = np.asarray(actions, dtype=np.float32)
        horizon = len(actions)
    state = np.zeros(4, dtype=np.float32)
    states = [state.copy()]
    used_actions = []
    for action in actions:
        x, y, vx, vy = state
        if context == 0:
            dvx = action[0] - 0.2 * vx
            dvy = action[1] - 0.2 * vy
        else:
            dvx = action[0] - 0.2 * vx + 1.2 * vy
            dvy = action[1] - 0.2 * vy - 1.2 * vx
        state = np.array(
            [x + DT * vx, y + DT * vy, vx + DT * dvx, vy + DT * dvy],
            dtype=np.float32,
        )
        states.append(state.copy())
        used_actions.append(action)
    return np.stack(states), np.stack(used_actions)


def make_dataset(n_episodes=400, horizon=EPISODE_HORIZON):
    S, A, C = [], [], []
    for i in range(n_episodes):
        ctx = i % 2
        s, a = rollout(ctx, horizon)
        S.append(s)
        A.append(a)
        C.append(ctx)
    return np.stack(S), np.stack(A), np.array(C)


S, A, C = make_dataset()
Out[14]:
Console
Trajectories: 400, transition horizon: 40, state dim: 4
Contexts present: [0, 1]

The two contexts start from rest at the origin. They initially agree; once velocity is nonzero, an informative state-action transition can distinguish their noiseless update laws. Longer windows add statistical evidence. For the visualization, each context in a pair receives the exact same action sequence, so differences within a pair isolate the context dynamics.

In[15]:
Code
trajectory_rng = np.random.default_rng(70)
paired_trajectories = {0: [], 1: []}
for _ in range(3):
    paired_actions = 0.5 * trajectory_rng.standard_normal(
        (EPISODE_HORIZON, 2)
    ).astype(np.float32)
    for context in (0, 1):
        states, _ = rollout(context, actions=paired_actions)
        paired_trajectories[context].append(states)
Out[16]:
Visualization
Three blue-orange trajectory pairs from the origin, paired by solid, dashed, or dotted style; orange context-B paths curve relative to blue context-A paths.
Three paired position rollouts under the two hidden contexts. Curves with the same line style share an action sequence; context B's rotational coupling curves the response and reduces average displacement across these three examples, although not in every pair.

Under the matched actions shown here, context B's rotational coupling curves the velocity response and reduces displacement on average across these three pairs. A state snapshot does not contain the context label, but after the dynamics have diverged its value may contain probabilistic evidence. A recent window supplies additional evidence about how velocity has evolved.

Preparing Windowed Transitions

For both window-conditioned models we build supervised examples of the form (window of recent states, current state, action, next state). The window offers evidence about which dynamics are in force, without supplying the context label or guaranteeing exact recovery.

In[17]:
Code
def make_transitions(S, A, contexts, window=WINDOW):
    """Construct (window_states, s_t, a_t, s_{t+1}) tuples for every episode."""
    N, T1, _ = S.shape
    T = T1 - 1
    wins, st, at, snext, ctxs = [], [], [], [], []
    for i in range(N):
        for t in range(window, T):
            wins.append(S[i, t - window : t].reshape(-1))
            st.append(S[i, t])
            at.append(A[i, t])
            snext.append(S[i, t + 1])
            ctxs.append(contexts[i])
    return (
        torch.tensor(np.stack(wins), dtype=torch.float32),
        torch.tensor(np.stack(st), dtype=torch.float32),
        torch.tensor(np.stack(at), dtype=torch.float32),
        torch.tensor(np.stack(snext), dtype=torch.float32),
        torch.tensor(np.stack(ctxs), dtype=torch.long),
    )


W_train, S_train, A_train, SN_train, C_train = make_transitions(
    S, A, C, window=WINDOW
)
Out[18]:
Console
Training transitions: 12800
Window flattened dim: 32

The output confirms the expected 32-feature dimension for an eight-state, four-dimensional window. The preceding slice and flattening operation supply the full contiguous recent-state window used for context inference.

Model Definitions

Three modules: a flat baseline that sees only (state, action), a history-conditioned model whose encoder compresses the window into a context latent, and a separate MoE with a learned router. The latter two receive more input and have more parameters than the flat baseline, so this is a diagnostic comparison, not a controlled estimate of an architectural advantage.

In[19]:
Code
class FlatDynamics(nn.Module):
    """Baseline: predict next state from (state, action) only."""

    def __init__(self, state_dim=4, action_dim=2, hidden=HIDDEN):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim + action_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, state_dim),
        )

    def forward(self, state, action):
        delta = self.net(torch.cat([state, action], dim=-1))
        return state + delta


class HistoryEncoder(nn.Module):
    """Compress a window of past states into a context latent."""

    def __init__(
        self, state_dim=4, window=WINDOW, latent_dim=LATENT_DIM, hidden=HIDDEN
    ):
        super().__init__()
        self.window = window
        self.net = nn.Sequential(
            nn.Linear(state_dim * window, hidden),
            nn.ReLU(),
            nn.Linear(hidden, latent_dim),
        )

    def forward(self, window_states):
        return self.net(window_states)


class FastWithContext(nn.Module):
    """Predictor conditioned on the inferred context latent."""

    def __init__(
        self, state_dim=4, action_dim=2, latent_dim=LATENT_DIM, hidden=HIDDEN
    ):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim + action_dim + latent_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, state_dim),
        )

    def forward(self, state, action, latent):
        delta = self.net(torch.cat([state, action, latent], dim=-1))
        return state + delta


class HistoryConditionedDynamics(nn.Module):
    """History encoder + context-conditioned one-step dynamics."""

    def __init__(
        self,
        state_dim=4,
        action_dim=2,
        window=WINDOW,
        latent_dim=LATENT_DIM,
        hidden=HIDDEN,
    ):
        super().__init__()
        self.history_encoder = HistoryEncoder(
            state_dim, window, latent_dim, hidden
        )
        self.fast = FastWithContext(state_dim, action_dim, latent_dim, hidden)

    def forward(self, window_states, state, action):
        latent = self.history_encoder(window_states)
        return self.fast(state, action, latent), latent

Training Loop

We share one training function across models. A lightweight forward adapter lets the flat, history-conditioned, and routed models use the same loop.

In[20]:
Code
def train_model(
    forward_fn,
    model,
    epochs=TRAIN_EPOCHS,
    lr=LEARNING_RATE,
    batch_size=BATCH_SIZE,
):
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    dataset = TensorDataset(W_train, S_train, A_train, SN_train)
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    losses = []
    for _ in range(epochs):
        epoch_loss, n = 0.0, 0
        for w, s, a, sn in loader:
            pred = forward_fn(model, w, s, a)
            loss = ((pred - sn) ** 2).mean()
            opt.zero_grad()
            loss.backward()
            opt.step()
            epoch_loss += loss.item() * w.shape[0]
            n += w.shape[0]
        losses.append(epoch_loss / n)
    return losses


def flat_forward(model, window, state, action):
    return model(state, action)


def history_forward(model, window, state, action):
    pred, _ = model(window, state, action)
    return pred
In[21]:
Code
flat_model = FlatDynamics()
flat_losses = train_model(flat_forward, flat_model)

history_model = HistoryConditionedDynamics()
history_losses = train_model(history_forward, history_model)

History-Conditioned vs. Flat Comparison

We compare training losses and multi-step rollout error on held-out test episodes.

In[22]:
Code
def rollout_error(model, forward_fn, S, A, window=WINDOW, horizon=EVAL_HORIZON):
    """Average Euclidean position error over recursive rollout steps."""
    errs = np.zeros(horizon)
    counts = np.zeros(horizon)
    for i in range(S.shape[0]):
        s_true = S[i, : horizon + window + 1].copy()
        s_pred = s_true[: window + 1].copy()
        for t in range(window, window + horizon):
            w = torch.tensor(
                s_pred[t - window : t].reshape(1, -1), dtype=torch.float32
            )
            s = torch.tensor(s_pred[t].reshape(1, -1), dtype=torch.float32)
            a = torch.tensor(A[i, t].reshape(1, -1), dtype=torch.float32)
            with torch.no_grad():
                s_pred_next = forward_fn(model, w, s, a).numpy()[0]
            s_pred = np.concatenate([s_pred, s_pred_next[None]], axis=0)
            errs[t - window] += np.linalg.norm(
                s_pred[t + 1, :2] - s_true[t + 1, :2]
            )
            counts[t - window] += 1
    return errs / counts


S_test, A_test, C_test = make_dataset(n_episodes=40, horizon=EPISODE_HORIZON)

Wait, there is a subtlety: make_dataset uses the shared RNG, so S_test differs from S. That is intentional: we want a fresh test set. We now visualize both training curves and multi-step rollout error.

Out[23]:
Visualization
Two decreasing training-loss curves; both approach values near zero, with the history-conditioned curve ending slightly lower in this run.
Training loss for the flat and history-conditioned models. Both reach very low loss; their final values should be compared directly rather than treated as a controlled architectural test.

Both training losses are very small in this run. The history-conditioned model ends slightly lower, but it also receives a state window and has more parameters. No decisive difference appears in recursive rollouts.

In[24]:
Code
flat_errs = rollout_error(flat_model, flat_forward, S_test, A_test)
history_errs = rollout_error(history_model, history_forward, S_test, A_test)
Out[25]:
Visualization
Flat and history-conditioned position-error curves rise together and remain close across the 20 rollout steps.
Mean Euclidean position error for the flat and history-conditioned models over a 20-step recursive rollout. The curves track closely, so this run does not show a meaningful history-conditioning advantage.

In this run the two models perform almost identically. The comparison neither isolates a benefit from history conditioning nor tests temporal abstraction: a controlled study would add a capacity-matched flat window baseline, match compute, and train across multiple seeds.

Adding a Mixture of Experts

Now let us compare the continuous context latent with a differentiable router and two experts. The experts see the same (state, action) inputs as the flat model, while the router uses the history window. The dense soft router blends both experts; it does not make a discrete mode decision, and specialization must be measured rather than assumed.

In[26]:
Code
class MoEDynamics(nn.Module):
    """Mixture of experts with a window-conditioned router."""

    def __init__(
        self,
        state_dim=4,
        action_dim=2,
        window=WINDOW,
        n_experts=N_EXPERTS,
        hidden=HIDDEN,
    ):
        super().__init__()
        self.n_experts = n_experts
        self.experts = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(state_dim + action_dim, hidden),
                    nn.ReLU(),
                    nn.Linear(hidden, hidden),
                    nn.ReLU(),
                    nn.Linear(hidden, state_dim),
                )
                for _ in range(n_experts)
            ]
        )
        self.router = nn.Sequential(
            nn.Linear(state_dim * window, hidden),
            nn.ReLU(),
            nn.Linear(hidden, n_experts),
        )

    def forward(self, window_states, state, action):
        x = torch.cat([state, action], dim=-1)
        gates = torch.softmax(self.router(window_states), dim=-1)
        deltas = torch.stack([e(x) for e in self.experts], dim=1)  # (B, K, D)
        delta = (gates.unsqueeze(-1) * deltas).sum(dim=1)
        return state + delta, gates


def moe_forward(model, window, state, action):
    pred, _ = model(window, state, action)
    return pred


moe_model = MoEDynamics()
moe_losses = train_model(moe_forward, moe_model)
In[27]:
Code
moe_errs = rollout_error(moe_model, moe_forward, S_test, A_test)
Out[28]:
Console
MoE final training loss: 0.0000
Flat final rollout error: 0.0261
History-conditioned final rollout error: 0.0267
MoE final rollout error: 0.0255

The MoE reaches a similarly small training loss, and its final position error is slightly below the flat baseline in this single seeded run. We next compute context-conditioned mean routing weights. This is one mode-alignment diagnostic, not proof of specialization: averages can hide sharp per-example routing or specialization along another axis.

In[29]:
Code
def routing_by_context(model, S, A, C, window=WINDOW):
    gates_by_context = {0: [], 1: []}
    for i in range(S.shape[0]):
        for t in range(window, S.shape[1] - 1):
            w = torch.tensor(
                S[i, t - window : t].reshape(1, -1), dtype=torch.float32
            )
            s = torch.tensor(S[i, t].reshape(1, -1), dtype=torch.float32)
            a = torch.tensor(A[i, t].reshape(1, -1), dtype=torch.float32)
            with torch.no_grad():
                _, gates = model(w, s, a)
            ctx = int(C[i])
            gates_by_context[ctx].append(gates.numpy()[0])
    arrays = {c: np.stack(values) for c, values in gates_by_context.items()}
    means = {c: values.mean(axis=0) for c, values in arrays.items()}
    mean_entropy = {
        c: float((-values * np.log(values + 1e-12)).sum(axis=1).mean())
        for c, values in arrays.items()
    }
    return means, mean_entropy


routing, routing_entropy = routing_by_context(moe_model, S_test, A_test, C_test)
context_routing_gap = 0.5 * np.abs(routing[0] - routing[1]).sum()
Out[30]:
Visualization
Grouped bars near one half for both experts and both contexts, showing little separation between context-conditioned mean gates.
Context-conditioned mean routing weights for two experts. Both contexts have near-uniform means, so this diagnostic shows little context alignment; the averages alone cannot rule out per-example specialization.

The context-conditioned means are nearly equal, so this diagnostic shows no clear one-to-one alignment between experts and hidden contexts. The conclusion is deliberately narrow. Per-example gate entropy is also useful, while expert-output differences, utilization, context probes, and expert ablations are needed before concluding that specialization is absent.

Out[31]:
Console
Context A mean gate entropy: 0.692
Context B mean gate entropy: 0.692
Total-variation gap between context mean gates: 0.004

Two observations close the example. First, neither the history-conditioned model nor the MoE establishes an advantage over the flat baseline in this single run. Second, the MoE's context-conditioned mean gates do not align with the known contexts. Mode alignment is useful when interpretable recovery is an explicit goal, but it should be paired with per-example routing statistics, probes, expert ablations, and task-level rollout performance. Low alignment alone does not prove that predictions average destructively across regimes.

Key Parameters

The key parameters for the worked history-conditioned and routed models are:

  • WINDOW: Number of past states supplied to the history encoder and router. More history can add context evidence when it is relevant, but also increases input dimension and shortcut risk.
  • LATENT_DIM: Dimensionality of the continuous context latent. The example does not claim that two dimensions are necessary or sufficient; that requires an ablation and a context probe.
  • N_EXPERTS: Number of experts. Two experts permit a one-to-one alignment with two contexts, but the alignment still has to be learned and evaluated.
  • HIDDEN: Width of the MLP hidden layers in the flat model, history encoder, context-conditioned predictor, router, and experts.
  • LEARNING_RATE / BATCH_SIZE / TRAIN_EPOCHS: Optimizer step size, minibatch size, and training passes in the shared loop.
  • DT: Forward-Euler simulation step. A smaller value generally reduces discretization error only when physical duration and the continuous action signal are held fixed while the step count changes accordingly.

Limitations and Impact

Hierarchical, hybrid, and modular architectures add failure modes beyond those of a single-level model. Hierarchy changes credit assignment: it introduces inter-level attribution and backup-design problems, while a well-chosen temporal abstraction can shorten long-horizon credit paths. Sparse slow updates create fewer slow computation nodes, but each node may aggregate gradients from many fast-step losses; signal strength depends on the objective and gradient paths. A fast predictor can also compensate for a drifting slow state unless training or evaluation checks the intended representation. Learned termination and update boundaries may fail only in closed-loop rollouts. These are risks of temporally sparse hierarchies, not mechanisms exercised by the worked example's per-step history encoder. Evaluate them with decision-centric and cross-level metrics in addition to one-step error.

Hybrid models inherit assumptions from their non-learned parts. If a simulator is structurally wrong, an additive residual may still correct it when the mismatch is observable and representable, but the task can become harder than learning a different model or boundary. Spurious predicates and missing axioms can produce invalid plans; execution monitoring, uncertainty, and abstraction repair help limit that risk. Differentiable hybrids add another hazard because a soft relaxation's gradient need not match the hard objective. A useful structural prior can improve efficiency when the learned component can represent the remaining task-relevant mismatch. Residual learning is therefore a conditional design choice for simulator hybrids, not a rule for every hybrid architecture.

Modular and omnimodal architectures face collapse, dominance, and shortcut risks. Aggressive sparse routing can concentrate traffic without an effective balancing design. Training and inference can diverge when they use different noise, temperature, truncation, capacity rules, or input distributions. In multimodal training, the objective may over-rely on whichever stream is easiest, whether vision, audio, touch, or another modality, and a shared latent may encode cross-modal shortcuts. A reconstruction score can therefore coexist with poor control utility. Use decision-aware objectives when task relevance is the problem, and combine them with balancing, masking, missing-input tolerance, uncertainty, or anti-shortcut interventions appropriate to the observed failure.

These architectures have delivered gains under specific conditions. Temporal abstraction can extend effective planning horizons; simulator hybrids can improve real-robot data efficiency when the simulator supplies a useful prior; broad multimodal architectures support training across mixed datasets and adaptation to deployment inputs; sparse MoE layers can add parameter capacity while evaluating few experts per example. None is automatic. Each exchanges a simpler architecture for additional assumptions, diagnostics, and failure modes.

Summary

This chapter studied four architectural families that lift limitations of flat world models.

  • Hierarchical models organize states or decisions across temporal, spatial, or abstraction levels. A slower update rate shortens the high-level decision chain; this can reduce compounding prediction error when coarse-transition error remains sufficiently small and fine-level error is controlled through accurate local dynamics or observation correction.
  • Hybrid models combine learned components with symbolic reasoning or hand-crafted simulators. Pipeline, differentiable, and residual patterns have distinct assumptions; a residual target is attractive when a useful simulator leaves an observable, representable mismatch.
  • Omnimodal models support broad input-output modality combinations. Cross-modal attention enables interaction. A task-dependent combination of reconstruction, cross-modal prediction, contrastive alignment, and decision-aware losses can encourage useful shared and complementary content, but none guarantees it alone. Masking and modality dropout can improve tolerance to missing inputs, which still depends on the missing streams and training coverage.
  • Modular models include mixture-of-experts architectures with learned routers. Sparse top-k routing can grow parameter capacity while evaluating few experts, but often needs an effective routing, capacity, or balancing mechanism to reduce harmful expert collapse; expert loads should be checked empirically.

The worked example compared a per-step history-conditioned context model and a separate routed model with a flat baseline. The window-conditioned models had more input and capacity, yet neither clearly improved multi-step position error in this single run. The MoE's context-conditioned mean gates showed little alignment with the two labels; richer routing diagnostics would be needed to characterize specialization.

The architectures in this chapter sit at the end of Part V. Part VI is organized around how world models are learned: predictive and generative objectives, interaction and passive-observation data, curation and synthesis, continual learning, and scaling. Part VI opens with predictive and generative objectives and examines how objective choice shapes learned representations and downstream use.

Quiz

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about hierarchical, hybrid, and omnimodal world models.

Comments

No comments yet. Be the first to share your thoughts!

Reference

Citation details

Cite or share this article.

BIBTEXAcademic
@misc{brenndoerfer2026hierarchicalhybrid, author = {Michael Brenndoerfer}, title = {Hierarchical, Hybrid, and Omnimodal Models}, year = {2026}, url = {https://mbrenndoerfer.com/writing/hierarchical-hybrid-omnimodal-world-models}, organization = {mbrenndoerfer.com}, note = {Accessed: 2026-10-15} }
APAAcademic
Michael Brenndoerfer (2026). Hierarchical, Hybrid, and Omnimodal Models. Retrieved from https://mbrenndoerfer.com/writing/hierarchical-hybrid-omnimodal-world-models
MLAAcademic
Michael Brenndoerfer. "Hierarchical, Hybrid, and Omnimodal Models." 2026. Web. October 15, 2026. <https://mbrenndoerfer.com/writing/hierarchical-hybrid-omnimodal-world-models>.
CHICAGOAcademic
Michael Brenndoerfer. "Hierarchical, Hybrid, and Omnimodal Models." Accessed October 15, 2026. https://mbrenndoerfer.com/writing/hierarchical-hybrid-omnimodal-world-models.
HARVARDAcademic
Michael Brenndoerfer (2026) 'Hierarchical, Hybrid, and Omnimodal Models'. Available at: https://mbrenndoerfer.com/writing/hierarchical-hybrid-omnimodal-world-models (Accessed: October 15, 2026).
SimpleBasic
Michael Brenndoerfer (2026). Hierarchical, Hybrid, and Omnimodal Models. https://mbrenndoerfer.com/writing/hierarchical-hybrid-omnimodal-world-models

About the author

Continue with the full handbook

This chapter is part of World Models Handbook. Use the handbook page to browse the complete table of contents and continue reading in sequence.

Explore World Models Handbook
Newsletter

Stay up to date

Get articles, book updates, and news delivered to your inbox.

No spam, unsubscribe anytime.

or

Join the community

Sign in to remove popups, track your reading progress, and join the discussion.