UniZero-style Transformer

• UniZero-style Transformer models are a class of world models for reinforcement learning that integrate a compact, history-aware transformer to predict latent dynamics and decision signals.
• They embed observations and actions into a unified latent space, enabling joint optimization of prediction and policy for efficient Monte Carlo Tree Search planning.
• By combining transformer-based latent modeling with MCTS, these models outperform traditional recurrent and convolutional architectures in multi-task and long-memory tasks.

UniZero-style Transformer models are a class of world models for reinforcement learning (RL) that integrate a modular transformer backbone to enable scalable planning, generalization across multiple tasks, and long-horizon credit assignment within a unified latent space. Inspired by value-equivalence principles underlying MuZero, UniZero departs from classical recurrent or purely convolutional world models by employing a small, history-aware transformer that jointly predicts latent dynamics and decision-oriented quantities. This design facilitates efficient Monte Carlo Tree Search (MCTS) in latent space, achieves superior scalability to diverse RL domains, and sets new state-of-the-art performance benchmarks on multi-task and long-memory tasks (Pu et al., 2024).

1. Architecture and Input Tokenization

At every timestep $t$, the agent receives an observation $o_t$ and accesses the preceding action $a_{t-1}$. The observation is embedded via a parameterized encoder $h_\theta$ (typically a small CNN followed by a linear layer and SimNorm normalization) to produce a latent state token: $z_t = h_\theta(o_t) \in \mathbb{R}^D,$ where $D$ is the transformer hidden dimension (e.g., $D=768$ for Atari). Actions are similarly embedded: $$e^a_t = \mathrm{Embed}_\theta(a_t) \in \mathbb{R}^D.$$ Sequences of $H$ consecutive frames/actions are concatenated into a $2H$-length token sequence

$o_t$0

with learnable position embeddings $o_t$1 added prior to transformer processing.

2. Transformer Configuration and Latent History Representation

The UniZero transformer backbone adopts a compact "nanoGPT"-style architecture:

• $o_t$2 transformer layers
• Hidden dimension $o_t$3 (Atari) or $o_t$4 (VisualMatch)
• Multi-head self-attention ($o_t$5 heads for Atari)
• Feedforward layers with GELU nonlinearity
• Dropout probability $o_t$6
• Learnable positional embeddings (no rotary/sinusoidal encoding)

After processing, output tokens are split into state and action hidden vectors $o_t$7. The entire trajectory representation is maintained in an implicit latent history, forming the memory substrate for both world modeling and policy decision heads.

3. Modular Output Heads: Latent Dynamics and Policy/Value

Two major modules are attached atop the transformer backbone:

Latent Dynamics Head ($o_t$8):

$o_t$9

where each "Linear" is a two-layer MLP (GELU, optional SimNorm).

Decision Head ($a_{t-1}$0):

$a_{t-1}$1

Here, $a_{t-1}$2 gives action logits (for $a_{t-1}$3) and $a_{t-1}$4 the predicted state value.

Both modules operate on the transformer-provided latent history, tightly coupling prediction of future game state (latent dynamics/reward) with decision making (policy/value).

4. End-to-End Training Objective and Optimization

The total objective function is a composite loss: $a_{t-1}$5 where:

• $a_{t-1}$6: soft target from an exponential moving average encoder,
• $a_{t-1}$7: $a_{t-1}$8-step TD return,
• $a_{t-1}$9: cross-entropy loss (for rewards, policy, value),
• $h_\theta$0: stop-gradient,
• Hyperparameters: $h_\theta$1 for Atari.

This design supports simultaneous optimization of prediction (dynamics/reward), planning (policy), and evaluation (state value) signals, balancing model-based and policy objectives.

5. Planning Implementation: Transformer-powered MCTS

At planning time, the transformer generates latent history used to feed MCTS, directly paralleling MuZero-style algorithms:

• Selection: Traverse from root latent $h_\theta$2 via PUCT rule using prior $h_\theta$3, value $h_\theta$4, and visit counts $h_\theta$5 as

$h_\theta$6

• Expansion: Compute $h_\theta$7.
• Backup: Propagate bootstrapped return

$h_\theta$8

and update $h_\theta$9, $z_t = h_\theta(o_t) \in \mathbb{R}^D,$0 accordingly.

• After $z_t = h_\theta(o_t) \in \mathbb{R}^D,$1 simulations, the improved policy is given by

$z_t = h_\theta(o_t) \in \mathbb{R}^D,$2

The transformer’s key-value cache underpins the entire search, capturing long-horizon dependencies with a fixed computation graph.

6. Scalability, Ablations, and Empirical Evaluation

UniZero demonstrates marked scalability and performance advantages:

• Outperforms MuZero-style and baseline world models in multitask settings and tasks requiring long-term memory (e.g., Atari multitask, VisualMatch).
• Maintains or exceeds state-of-the-art results in single-task settings (Atari, DMControl).

Key hyperparameters and ablations reveal:

• SimNorm normalization on latent representations provides enhanced stability compared to Softmax or Sigmoid.
• EMA soft targets are crucial for convergence; hard targets or no targets induce instability.
• Inference context $z_t = h_\theta(o_t) \in \mathbb{R}^D,$3 suffices for standard Atari; longer context is beneficial in long-dependency regimes.
• Ablations indicate that image reconstruction decoders do not add value and are excluded from the final model.
• The architecture allows for efficient long-horizon planning with reduced compute by unifying model and policy updates within a single transformer.

7. Significance and Distinctive Innovations

UniZero’s design addresses fundamental bottlenecks in MuZero-style world models, notably by:

• Leveraging a transformer backbone for modular, history-aware latent modeling.
• Producing both predictive (dynamics, reward) and decision-oriented (policy, value) signals from a joint latent space.
• Enabling joint optimization and planning, which are critical for heterogeneous domains and tasks with temporal correlations spanning hundreds of steps.

A plausible implication is that transformer-based latent world models of this form may become foundational modules for future generalist RL agents, capable of efficient cross-task transfer and reasoning over long temporal horizons (Pu et al., 2024).