EfficientZero V2: Sample-Efficient RL

• The paper introduces EfficientZero V2, a model-based RL framework that extends to both discrete and continuous actions while improving sample efficiency and generalization.
• It leverages sampling-based Gumbel search, search-based value estimation, and mixed value targeting to deliver robust performance under data constraints.
• Empirical results and ablations show EZ-V2 outperforms SOTA baselines like DreamerV3 on benchmarks such as Atari and DeepMind Control Suite.

EfficientZero V2 (EZ-V2) is a general-purpose, sample-efficient model-based reinforcement learning (RL) framework that extends the EfficientZero family to discrete and continuous action spaces as well as high- and low-dimensional observational domains. EZ-V2 integrates innovative planning, value estimation, and architecture enhancements to provide robust performance and sample efficiency across diverse benchmarks, consistently outperforming state-of-the-art (SOTA) baselines including DreamerV3 and earlier EfficientZero variants under data-constrained regimes (Wang et al., 2024).

1. Background and Prior EfficientZero Developments

The precursor EfficientZero (V1) [Ye et al., 2021] is a model-based RL algorithm that constructs a latent predictive environment model and performs planning with Monte Carlo Tree Search (MCTS). Its main computational graph comprises:

• Representation Network $H(o_t)\rightarrow s_t$: Encodes the current observation (e.g., raw image) into a latent state.
• Dynamics Network $G(s_t,a_t)\rightarrow (s_{t+1},r_t)$: Predicts the next latent state and immediate reward from the current latent state and action.
• Policy and Value Heads: $P(s_t)$ outputs a policy $\pi_t(a|s_t)$; $V(s_t)$ predicts the state value $v_t$.
• Planning/Target Extraction via MCTS: Utilizes the latent model $\{H,G,P,V\}$ for simulated rollouts, Q-value estimation, and deriving a target policy $\pi^t$ and value $z^t$ for supervised training.

The loss per timestep is a composite:

$$\mathcal{L}_t = \mathcal{L}_r(u_t, \hat{r}_t) + \mathcal{L}_p(\pi^t, \pi_t) + \mathcal{L}_v(z^t, v_t) + \mathcal{L}_g(s_{t+1}, \hat{s}_{t+1})$$

where $G(s_t,a_t)\rightarrow (s_{t+1},r_t)$0 is a temporal-consistency loss enforcing next-state alignment.

2. Key Innovations in EfficientZero V2

EZ-V2 introduces significant methodological and architectural advances over its predecessor, designed to improve sample efficiency, generalization, and versatility:

2.1 Sampling-based Gumbel Search for Continuous Control

Instead of discrete MCTS, EZ-V2 samples $G(s_t,a_t)\rightarrow (s_{t+1},r_t)$1 candidate actions at the root node:

• $G(s_t,a_t)\rightarrow (s_{t+1},r_t)$2: $G(s_t,a_t)\rightarrow (s_{t+1},r_t)$3 samples from the current policy $G(s_t,a_t)\rightarrow (s_{t+1},r_t)$4,
• $G(s_t,a_t)\rightarrow (s_{t+1},r_t)$5: $G(s_t,a_t)\rightarrow (s_{t+1},r_t)$6 samples from a broader prior $G(s_t,a_t)\rightarrow (s_{t+1},r_t)$7 (e.g., flattened $G(s_t,a_t)\rightarrow (s_{t+1},r_t)$8),
• $G(s_t,a_t)\rightarrow (s_{t+1},r_t)$9 total actions, with the Gumbel-Top-$P(s_t)$0 trick for sampling without replacement [Kool et al., 2019].

Actions are selected using a sequential halving bandit algorithm [Karnin et al., 2013] to efficiently compare $P(s_t)$1 estimates and expand the tree. At non-root nodes, samples are only from $P(s_t)$2, reducing variance and increasing search depth.

A policy-improvement guarantee is established: the selected action $P(s_t)$3 from $P(s_t)$4 satisfies

$P(s_t)$5

2.2 Search-Based Value Estimation (SVE)

SVE addresses off-policy errors in replayed transitions by re-estimating the value of any state $P(s_t)$6 via $P(s_t)$7 repeated tree-searches:

$P(s_t)$8

An explicit mean squared error bound is provided (Theorem 4.3):

$P(s_t)$9

which vanishes as the model error $\pi_t(a|s_t)$0 approaches zero or as search depth increases.

2.3 Mixed Value Targets

Targets for value learning are dynamically selected: recent transitions employ multi-step TD, whereas heavily replayed transitions adopt SVE for reduced off-policy bias:

$\pi_t(a|s_t)$1

where $\pi_t(a|s_t)$2 is the buffer index, $\pi_t(a|s_t)$3 the buffer size, and $\pi_t(a|s_t)$4 are thresholds.

2.4 Action Embedding Layer

Continuous-space actions are projected into a 64-dimensional embedding via Linear→LayerNorm→ReLU, enhancing the expressivity and generalization of the latent dynamics model.

2.5 Priority Pre-Calculation

Transition priorities upon arrival are initialized by Bellman errors under the current model, rather than set to max-priority, thus accelerating effective data reuse during early training.

2.6 Architecture Variations

• Visual inputs adopt convolutional and residual towers from V1.
• One-dimensional/proprioceptive inputs use a linear+LayerNorm+Tanh preprocessing followed by three Pre-LN Transformer-style residual blocks (size: 256, output: 128, ReLU).

3. Algorithmic Workflow

The end-to-end EZ-V2 procedure involves alternating data collection and model-based offline reanalysis as described in the algorithm pseudocode:

$\pi_t(a|s_t)$6

4. Network Architectures and Hyperparameters

EZ-V2 network and optimization pipeline features:

Component	Architecture	Purpose
Representation H	1-D: Linear→LN→Tanh→3×[Pre-LN Res(256)]→128; Image: Conv + Residual towers	Latent state encoding
Dynamics G	$\pi_t(a|s_t)$5 (128-D) ∥ Action emb. (FC→64→LN→ReLU); 3×[Pre-LN Res]	State/reward prediction
Policy/Value Heads	FC→LN→ReLU→MLP(256); Policy: μ,σ (5·Tanh/Softplus); Value: categorical	Policy/value prediction

Key training parameters include:

• Replay buffer: 10⁶ capacity
• Batch size: 256
• Discount: γ = 0.997
• UTD ratio: 16
• Unroll steps: 5
• Simulation count: 32 (16 for Atari)
• Sampled actions: 16 (8 for Atari)
• Priority exponent: 0.9
• Loss weights: αᵣ=1.0, αₚ=1.0, αᵥ=0.25, αₛ=2.0, entropy=5e-3
• Optimizers: Adam (proprio, lr=3e-4, wd=2e-5); SGD (vision & Atari, lr=0.2, m=0.9, wd=1e-4).

5. Empirical Results and Comparative Performance

EZ-V2 is evaluated across canonical RL benchmarks under stringent data limits:

• Atari 100k: Mean normalized score 2.428, median 1.286. EfficientZero V1: mean 1.945, DreamerV3: mean 1.120. EZ-V2 achieves superhuman performance in 23/26 games within a two-hour training budget.
• DeepMind Control Suite:
  • Proprioceptive (50k–100k steps): EZ-V2 mean 723.2, median 800.4; outperforming DreamerV3 (mean 517.1, median 543.4).
  • Vision Control (100k–200k steps): EZ-V2 mean 726.1, median 788.1; DreamerV3 mean 410.3, median 330.6.
• Total Benchmark Coverage: Outperforms or matches SOTA in 50 of 66 tasks.

This demonstrates effective use of planning, tailored targets, and representation strategies to yield robust policy learning under limited data.

6. Ablation Studies and Analysis

Two primary ablation axes are analyzed:

• Gumbel Search vs. Sample MCTS: EZ-V2’s Gumbel-based planner with Nsim=32 consistently surpasses Sample MCTS (Nsim=50) in fidelity and convergence speed, even with as few as 8 simulations on hard tasks such as 'Quadruped Walk' and 'Acrobot Swingup'.
• Value Targeting: Mixed value targets provide faster and more stable learning than pure multi-step TD or Double-Q Bellman methods, especially under sparse reward regimes.

The combination of sampling-based search and SVE with mixed targets is shown empirically to mitigate off-policy bias and enhance value estimation on replayed transitions. A plausible implication is that this approach generalizes better to environments with sparse data or infrequent extrinsic rewards.

7. Sample Efficiency and Generalization Capabilities

The observed sample efficiency results from:

• Executing only tens of model-based simulations per planning step versus the thousands typical in online MPC.
• Self-supervised temporal consistency losses, which improve representation learning.
• Enhanced data reuse via priority warm-up and offline value reanalysis (SVE).

EZ-V2 utilizes a unified set of architectures and hyperparameters for tasks with discrete or continuous action spaces, both visual and proprioceptive input modalities, and for environments with dense or sparse reward signals. Theoretical analysis (Theorem 4.3) guarantees SVE error decays with improved model quality or deeper search, while empirical results indicate consistent outperformance versus previous SOTA under tight step budgets (Wang et al., 2024).