NestedVLA: Consolidate RL Skills with Hypernetworks

Suppose you have three robot policies fine-tuned from the same starting checkpoint. One works well on MetaWorld, another on ManiSkill, and a third on CALVIN. Keeping all three full models duplicates most of their parameters. Averaging their weights saves storage, but updates that help one environment can disrupt another.

NestedVLA learns a hypernetwork that generates an action-expert update for the selected task. It preserves a shared base while retaining the information needed to materialize different policies. This tutorial explains the method, checks what the released implementation actually supports, and walks through a small PyTorch experiment. Robot results come from the paper; the synthetic example below does not reproduce those benchmarks.

The primary sources are NestedVLA: Learning to Consolidate and Generate Skills for Vision-Language-Action Models by Guodong Du and Wanyu Lin, NeurIPS 2026, the v3 paper, and the official repository. As checked on October 26, 2026, the repository releases the core algorithm without benchmark datasets, simulators, training checkpoints, or experiment launch scripts.

1. The concepts you need first

A Vision-Language-Action model turns visual observations and language instructions into robot actions. In the π0.5 backbone used by the paper, the vision-language model processes multimodal information, while an action expert produces continuous actions. For background, read how VLA models connect images, language, and control.

Reinforcement learning improves a policy through interactions and rewards. In this study, RL creates the specialized source policies. Consolidation happens afterward: the hypernetwork learns their parameter updates through regression. Training the generator is therefore a different optimization problem from training a robot policy inside a simulator.

A hypernetwork produces parameters for another network. A normal robot policy maps observations to actions. The NestedVLA generator maps a task identifier and patch addresses to weight updates. Those updates are assembled into a policy, which then receives camera images, robot state, and instructions through the original VLA interface.

Think of each skill as an edit to the same instruction manual. Weight averaging tries to produce one edit that works for everyone. NestedVLA learns to reconstruct the appropriate edit when a skill is selected. The evaluated system still needs a task identity; generating weights does not establish that it can invent an unseen skill from an arbitrary instruction.

2. Store and generate updates around a common base

Let Theta_base be the original checkpoint and Theta_m the policy fine-tuned for task m. Their difference is a task vector:

tau_m = Theta_m - Theta_base

action_m = action_base + generated_action_delta(m)
VLM_m    = VLM_base    + stored_compressed_VLM_delta(m)

All policies must originate from the same base and have compatible parameter names, shapes, and meanings. Matching tensor dimensions alone is insufficient if the checkpoints come from different backbone revisions.

NestedVLA separates the update into two branches. The action-expert delta is learned by a shared generator. The VLM delta is compressed and retained separately for each task. Consequently, a single hypernetwork does not replace every task-specific artifact: compressed VLM residuals, masks, and task-specific generator components remain part of the system.

Source: Figure 1 of the NestedVLA paper. The diagram contrasts merged weights with generated updates around a shared base.

The separate VLM branch matters. In the Appendix F sanity test, the authors keep the ground-truth action expert for each task but merge the VLM. CALVIN still drops to zero chain completion for the tested merging methods. A correct action expert cannot compensate for every change to the representations it receives.

This is post-hoc consolidation: the source policies have already learned their behavior. The method does not require pooling their original training data during generator fitting. It does require access to the learned checkpoints and a compatible base, and meaningful validation still requires the relevant evaluation environments.

3. Compress deltas and tokenize parameter patches

Dense updates for a model with billions of parameters are large regression targets. NestedVLA first keeps high-magnitude entries using a top-K support mask, then applies a rescale factor. The main paper configuration retains 10% per module and uses a global factor of 1.1.

compressed_delta = alpha * delta * mask
mask[i] = 1 for retained positions
mask[i] = 0 elsewhere

The support is computed once and frozen. The generator learns the retained values rather than simultaneously discovering where an update should exist. At inference, its output is masked again, so values outside the selected support never enter the reconstructed action expert.

Next, each parameter tensor is represented as a matrix and divided into square tiles. The paper uses 256 × 256 patches. Small tensors and incomplete edge tiles are padded; decoding removes that padding and restores the original shape.

Each patch has four address fields: parameter type, transformer layer index, patch row, and patch column. These are weight patches, not camera-image patches. Two tiles with the same coordinates can belong to different attention or MLP tensors, so their parameter type and layer must also be represented.

The public DeltaTokenizer implements encoding and decoding. Its patch_scales() computes an RMS scale for each patch with a lower bound to avoid division by zero. Standardization makes target magnitudes easier to optimize; destandardization restores the original units. Keep separate scales for every task. Reusing one task's scales for another silently changes the generated update even when tensor shapes still match.

Padding also deserves attention when experimenting with very small tensors. A large patch around a tiny matrix mostly contains artificial zeros. Start with dimensions divisible by the patch size so you can understand the core behavior before investigating padding and sparsity interactions.

4. The paper architecture versus the public core

In the research architecture, a task embedding passes through a small Conv2D encoder and global average pooling. The resulting task feature is broadcast across patches and concatenated with patch-address embeddings. A shared fully connected trunk, task-specific LoRA path, and dedicated output heads produce the weight tiles.

Source: Figure 3 of the NestedVLA paper. The released core is simpler than this research architecture.

Checking model.py, compression.py, and the paper reveals several material differences:

These differences prevent treating the README example as a reproduction of the headline result. The core provides a useful tensor-level API and an understandable starting point. Reproducing the complete system requires additional architecture, integration, checkpoints, and experimental configuration.

There is also an optimization detail: even with its outside-mask weight set to zero, the public loss averages over every tensor position. The main paper equation normalizes by the number of retained entries. With a fixed keep ratio these objectives are related by a scale factor, but changing that ratio affects gradient magnitude and can require retuning optimization settings.

5. Install a small experimental environment

The package metadata requires Python 3.10 or newer and PyTorch 2.0 or newer. Use an isolated environment. For GPU experiments, install a PyTorch build compatible with your CUDA setup; a CPU is sufficient for the small example here.

The following commands are for your own practice environment. Installing from a source archive avoids requiring a Git checkout:

python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install torch
python -m pip install "https://github.com/WanyuGroup/NeurIPS2026-NestedVLA/archive/refs/heads/main.zip"
python -c "from nested_vla import DeltaTokenizer, NestedVLAGenerator; print('Import OK')"

The main branch can change after publication. For a reproducible experiment, record the exact source revision and dependency versions that you install. A command pointing to a moving branch is convenient for exploration, but it does not identify a permanent implementation.

Begin with patch size 4 rather than copying the paper's 256 immediately. The next example uses matrices divisible by 4, making tokenization and the selected support easy to inspect. Its smaller hidden dimensions and higher keep ratio are teaching choices, not experimental settings from the paper.

6. Train a delta generator with the public API

Run the complete block in one Python session. It creates three synthetic task updates around a shared base. No robot trajectories are loaded. Setting out_of_mask_weight=0.0 restricts learning to retained positions, while keeping the public API's averaging convention described above.

import torch
from nested_vla import (
    DeltaTokenizer, NestedVLAGenerator, make_topk_mask,
)

torch.manual_seed(7)
torch.set_num_threads(1)

# Every task must share parameter names, shapes, and ordering.
base = {
    "layers.0.attn.weight": torch.randn(8, 8),
    "layers.1.mlp.weight": torch.randn(8, 12),
}
deltas = [
    {name: 0.02 * torch.randn_like(value)
     for name, value in base.items()}
    for _ in range(3)
]

tokenizer = DeltaTokenizer.from_state_dict(deltas[0], patch_size=4)
encoded = [tokenizer.encode(delta) for delta in deltas]
for name, value in tokenizer.decode(encoded[0]).items():
    assert torch.equal(value, deltas[0][name])

masks = torch.stack([
    make_topk_mask(patches, keep_ratio=0.25)
    for patches in encoded
])
scales = torch.stack([
    tokenizer.patch_scales(patches) for patches in encoded
])
targets = torch.stack([
    tokenizer.standardize(patches, scale)
    for patches, scale in zip(encoded, scales)
])

model = NestedVLAGenerator(
    tokenizer.metadata, num_tasks=3,
    task_dim=16, slot_dim=8, hidden_dim=64,
    grid_channels=4, grid_size=4,
    out_of_mask_weight=0.0,
)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3)
task_ids = torch.arange(3, dtype=torch.long)

model.train()
for step in range(500):
    optimizer.zero_grad(set_to_none=True)
    prediction = model(task_ids)
    loss = model.masked_loss(prediction, targets, masks)
    loss.backward()
    optimizer.step()
    if step % 100 == 0:
        print(step, float(loss.detach()))

model.eval()
generated = model.generate_delta(
    task_id=1, tokenizer=tokenizer,
    mask=masks[1], scales=scales[1],
)
reference = tokenizer.decode(encoded[1] * masks[1])
error = torch.sqrt(torch.stack([
    (generated[name] - reference[name]).square().mean()
    for name in reference
]).mean())
print("Sparse-target RMSE:", float(error))

# Rebuild from the original base rather than accumulating updates.
materialized = {
    name: base[name] + generated[name] for name in base
}
for name in base:
    assert materialized[name].shape == base[name].shape
print("Shapes:", {name: tuple(x.shape)
                  for name, x in materialized.items()})

First check that imports work, tokenization round-trips exactly, loss is finite, and reconstructed shapes match. Optimization should generally reduce the small regression objective, but no executed training log or measured convergence threshold is claimed here. Actual values depend on the installed implementation and optimization behavior.

The demo retains four entries from each 4 × 4 patch because its keep ratio is 0.25. Its reference is the masked target, not the dense original delta. Comparing against the dense update would combine two different errors: information removed by compression and approximation error introduced by the generator.

The printed RMSE averages the per-tensor mean squared errors before taking a square root. It is a small reconstruction diagnostic, not a benchmark success metric. Even excellent parameter reconstruction can leave sensitive control decisions wrong, so the final test must be closed-loop behavior.

7. Replace synthetic updates with real checkpoints

For a real experiment, replace random delta creation with a parameter-wise subtraction between each fine-tuned checkpoint and the original base. Select only the action-expert keys identified by your framework. Validate keys and shapes, then perform subtraction in float32 to reduce precision loss when two weights are close.

Do not guess an action-expert prefix from another repository. Inspect the actual model structure. Frameworks can use different names and state-dict layouts; buffers and non-floating entries also need deliberate handling. Preserve a stable key order across tasks and verify the common base using experiment metadata.

The VLM branch needs its own extraction, compression, and storage path. The public core does not include a complete π0.5 loader or automatic separation of all VLM and action-expert parameters. Omitting VLM residuals changes the evaluated system, even if your action-side code matches the API perfectly.

For context on collecting the source skills, see the VLA reinforcement-learning training workflow. The paper independently fine-tunes π0.5 with FlowSDE and uses RLinf for evaluation. Generator training does not replace that acquisition stage or make its compute cost disappear.

The paper's consolidation recipe uses 6,000 steps, batch size 3, base learning rate 3e-4, 600 warmup steps, and a 1,500-step cosine tail. Task weights are 1.0, 0.4, and 1.0 for MetaWorld, ManiSkill, and CALVIN. The reported run uses one A6000 with 48 GB, about ten hours of training, and 42.6 GB peak GPU memory. These are consolidation costs, not the total cost of producing the original RL policies.

8. Inference means materializing a selected policy

At a task switch, query the generator, restore the action delta using that task's scales and mask, retrieve the corresponding compressed VLM residual, and add both updates to the original base. The resulting policy then runs with the framework's image preprocessing, language tokenizer, state normalization, and action decoder.

There is no need to regenerate hundreds of millions of weights at every control step when the selected task remains unchanged. Caching materialized policies can reduce switching work, but its memory cost must be counted. Always rebuild from a clean base: adding task B's delta to weights already containing task A's delta creates an unintended policy.

Define your rescale convention explicitly. The current generate_delta() only destandardizes and applies the mask; it does not automatically multiply by 1.1. If targets already include alpha, applying it again after decoding scales the update twice. The tutorial keeps alpha at one to make the minimal API behavior unambiguous.

Before deployment, the usual policy validation and deployment checks still apply. Correct tensor shapes do not verify camera coordinates, action units, or actuator limits. Evaluate the reconstructed policy in the appropriate simulator before considering physical execution.

9. What the reported results establish

The following rows come from Table 1 in the v3 paper. Scores are percentages, and storage is relative to one full checkpoint.

The 41.1 improvement is in percentage points over Free-Merging. It is neither a latency reduction nor a claim that inference is 41.1% faster. NestedVLA remains 13.8 points below separate fine-tuned checkpoints on the reported average. ManiSkill is only 0.9 points behind, while the MetaWorld gap is still 34.4 points.

MetaWorld and ManiSkill use success_once: reaching the goal at any time in a rollout counts as success. CALVIN uses success_at_end for completing an entire five-subtask chain. The three-column average is the paper's summary across these benchmarks; it is not a pooled success rate over trajectories sharing one identical definition. Evaluation uses three seeds, with a typical standard deviation around 2.7%.

Table 3 helps distinguish compression limits from generator limits. Even the actual ground-truth delta, masked at 10% without learning a generator, reaches only 48.4% on MetaWorld. Raising the keep ratio to 20% increases that ideal reconstruction to 60.2%. A larger hypernetwork cannot recover target information that the compression stage has already discarded.

The ablations also discourage treating every architectural component as essential. Removing the LoRA path produces lower reconstruction losses at the 2,000-step comparison point. That is a specific optimization result, not evidence that every VLA application should abandon LoRA. Similarly, the headline 1.5× storage applies to three source policies; residuals, masks, and task-specific generator components continue growing as more tasks are added.

10. A practical sequence for your first experiment

Work through three levels. First, use small tensors to verify encoding, decoding, masking, and task-conditioned reconstruction. Second, use real checkpoints but evaluate compressed ground-truth deltas before fitting a generator. This measures the behavior preserved by your sparsity budget. Third, substitute generated deltas and evaluate with the same environments, rollout seeds, normalization, and success metrics.

If the second level fails, investigate compression before extending training. If compressed checkpoints work but generated checkpoints fail, inspect scales, masks, parameter ordering, and generator capacity. If reconstruction looks good but rollouts remain poor, inspect preprocessing, the retained VLM branch, and control sensitivity.

NestedVLA offers a concrete approach to reusing independently acquired skills through conditional weight generation. Its usefulness in your stack depends on the available checkpoints and integration work. The public tensor-level release is a good place to understand the mechanism; closed-loop evidence is what determines whether the reconstructed robot policy actually retains its skills.

Related Posts

• OpenVLA: understanding Vision-Language-Action architecture
• SimpleVLA-RL: the training workflow
• VLA and LeRobot: PEFT, validation, and policy deployment