- Muse-Robotics-1
- Inputs and Observation Model
- Vision Encoding
- Language Encoding
- Fusion and the Perceive Stack
- Deliberation Through Latent Plan Tokens
- Action Generation With Rectified Flow
- Training and Optimization
- Using the Model
- Configuration
- Demo
- Limitations
- Citation
- License
- Image Preprocessing and Augmentation
- Action Representation
- Evaluation
- Reproducibility
- Adapting the Model
- Inputs and Observation Model
Muse-Robotics-1
Muse-Robotics-1 is a compact, fully trainable vision-language-action (VLA) model for language-conditioned robotic manipulation. It takes two camera views (an exterior third-person view and a wrist-mounted view), the robot's proprioceptive state, and a free-form natural-language instruction, and outputs a short horizon of continuous actions. The defining idea is that the policy does not merely react: it maintains a small set of learned latent "plan" tokens that are refined over several internal steps before any action is produced, giving the network a bounded, differentiable space in which to deliberate about the task.
The model is built from scratch rather than relying on frozen pretrained backbones. Every parameter — vision, language, fusion, deliberation and action generation — is trained jointly end to end. This keeps the system small (roughly 120 million parameters) and self-contained, while ensuring that all representations are grounded directly in the control objective instead of borrowed from unrelated tasks.
Inputs and Observation Model
The policy consumes a fixed observation bundle at each control step. The two images are 224×224 RGB frames captured from distinct vantage points. The proprioceptive state is a low-dimensional vector covering joint or end-effector positions and an encoding of the gripper. The instruction is an arbitrary string, which may be empty; empty instructions are handled explicitly rather than treated as an error.
Normalization is critical for stability. The mean and standard deviation for both state and action are computed over the training distribution and embedded directly in the published configuration file. At inference the policy denormalizes internally, so the consumer of the model receives commands already in the native robot space and never needs access to the original dataset statistics.
Vision Encoding
Each camera frame is split into a grid of 16×16 pixel patches, linearly embedded into tokens, and augmented with sinusoidal position encodings that preserve the original two-dimensional layout. The exterior and wrist views are processed by a single shared Vision Transformer with a small (S-scale) architecture. Weight sharing between views encourages the network to learn viewpoint-invariant features such as object identity and spatial relations, while a learned view-type embedding lets the trunk still tell the two cameras apart.
The vision encoder is intentionally shallow. Because the whole model is trained from random initialization on embodied data, a deep pretrained backbone would be both unnecessary and a source of mismatch; a compact encoder keeps training tractable on modest hardware while still providing rich, task-relevant visual tokens.
Language Encoding
Instructions are encoded with a byte-level byte-pair-encoding tokenizer of size 8192, with dedicated symbols for padding, sequence boundaries and the empty instruction. The token sequence is passed through a six-layer bidirectional transformer that produces contextual embeddings for each word piece. Attention is correctly masked so that padding never influences the representation, and positional encodings preserve word order.
Training the language encoder jointly with the control loss means its representations become tightly coupled to actionable meaning. The encoder is width-matched to the vision encoder so that their outputs can be concatenated into a single sequence without an additional projection bottleneck that would otherwise narrow the information flow.
Fusion and the Perceive Stack
After encoding, the visual tokens, language tokens and the projected state token are concatenated into one heterogeneous sequence. Each token receives a learned type embedding (indicating whether it came from vision, language or proprioception) and a learned spatial embedding that complements the sinusoidal codes. This lets the model treat all modalities homogeneously while never losing track of where each piece of information originated.
The combined sequence is processed by a twelve-layer "perceive" transformer stack. These layers use pre-normalization and feed-forward expansion of factor four, and they build a shared memory that captures cross-modal dependencies — for example, binding the word "drawer" to the corresponding region in the image. The output of this stack is the context that the deliberation stage will later attend to.
Deliberation Through Latent Plan Tokens
Rather than decoding actions directly from the perceive memory, Muse-Robotics-1 first distills the relevant information into eight learned latent plan tokens. These tokens are initialized as parameters and have no direct correspondence to any input; they are a learned bottleneck through which all control decisions must pass.
The plan tokens are refined through a "thinking" stage consisting of three iterations of a shared transformer block. The block applies self-attention among the plan tokens, then cross-attention from the plan tokens to the full perceive memory, then a feed-forward update — all wrapped in residual connections. A distinct learned iteration embedding is added on each pass, so the single shared block behaves differently depending on whether it is the first, second or third thinking step. This gives the network a form of internal reasoning: across iterations the plan tokens become increasingly specific and actionable, while the parameter count stays low because the block is reused.
The final pooled plan representation is the sole conditioning signal for the action generator. Because the plan is a continuous, differentiable vector, the entire deliberation process is trained jointly with the control objective and can be inspected or ablated without changing the surrounding code.
Action Generation With Rectified Flow
The action generator is a transformer that operates over a chunk of sixteen future actions. Each action is a seven-dimensional delta command covering position, rotation and the gripper. Generation follows a rectified-flow formulation: at training time a clean action chunk is linearly interpolated toward Gaussian noise, and the network is trained to predict the velocity field that would move noise back to the clean sample. The loss is the squared error between predicted and target velocity, averaged over the horizon and over the active action dimensions.
During inference the model integrates this learned velocity field starting from random noise, using ten Euler steps. The plan tokens condition the denoising process by being concatenated into the action sequence and attended to jointly, so the generated motion is guided by the latent intention without a separate cross-attention branch. The continuous outputs are denormalized with the embedded statistics before being returned.
A small auxiliary objective runs in parallel: a pooled plan vector is asked to predict future proprioceptive states a few steps ahead. This gently shapes the latent to be forward-looking. Its weight is kept low so that it supports, rather than competes with, the primary flow objective.
Training and Optimization
The model is trained with AdamW under a cosine learning-rate schedule with warmup. Gradient clipping and an exponential moving average of the weights are used; the EMA copy is what is exported for inference because it typically yields smoother, more consistent control. Mixed-precision (bfloat16) training keeps memory and compute requirements modest.
Because all parameters are learned from scratch, careful initialization and deterministic seeding of data order, augmentation and weights are used so that runs are reproducible and can be resumed exactly. Checkpoints store the model, optimizer, scheduler and EMA state, with a stable pointer to the latest step and a separate safetensors file containing only the inference weights.
Using the Model
The published checkpoint is self-contained. Loading it requires only the configuration, the tokenizer and the safetensors weights:
from nova_vla.config import load_config
from nova_vla.model import NovaVLA
from nova_vla.tokenizer_util import load_tokenizer
from safetensors.torch import load_file
cfg = load_config("config.json")
model = NovaVLA(cfg.model)
model.load_state_dict(load_file("model.safetensors"))
tokenizer = load_tokenizer("tokenizer.json")
For deployment, a FastAPI policy wraps the model and accepts base64-encoded frames together with the current state and instruction, returning denormalized action chunks:
from nova_vla.serve import Policy
policy = Policy("hf_muse-robotics-1", device="cuda")
out = policy.act(front_b64, wrist_b64, state, "close the drawer")
The serving path shares its image preprocessing and tokenization with training, which removes a common source of train-to-deploy skew.
Configuration
The architecture is fully described by the embedded configuration. The notable choices are: image size 224 with patch size 16; vision and language encoders at width 384 with depths 12 and 6 respectively; a trunk and action expert at width 512; a twelve-layer perceive stack; three thinking steps over eight plan tokens; and an action chunk of sixteen. These values can be overridden through dot-path notation on the configuration dataclasses, allowing architecture sweeps without code changes.
Demo
The clip below shows Muse-Robotics-1 running on held-out episodes from a real robot dataset. For each timestep the policy receives the two camera views and proprioceptive state, then predicts a sixteen-step action chunk. The overlay compares the predicted actions (cyan) against the dataset ground truth (yellow) and reports the per-step error.
Limitations
Muse-Robotics-1 is best understood as a research base for adaptation rather than a universal controller. It will perform best on tasks close to its training distribution and can degrade on long-horizon goals, contact-rich interactions or substantially novel scenes. The model performs no explicit uncertainty estimation, collision avoidance or high-level planning, so any real deployment should surround it with workspace limits, velocity and force thresholds, and appropriate human oversight.
Citation
@article{khazatsky2024droid,
title={DROID: A Large-Scale In-The-Wild Robot Manipulation Dataset},
author={Khazatsky, Alexander and others},
journal={arXiv:2403.12945},
year={2024}
}
@inproceedings{walke2023bridgedata,
title={BridgeData V2},
author={Walke, Homer and others},
booktitle={CoRL},
year={2023}
}
@article{black2024pi0,
title={pi0: A Vision-Language-Action Flow Model},
author={Black, Kevin and others},
journal={arXiv:2410.24164},
year={2024}
}
License
MIT
Image Preprocessing and Augmentation
During training, frames are randomly resized and cropped before being normalized to a symmetric range, and photometric jitter is applied to improve robustness to lighting and camera variation. At evaluation and inference the same resize and normalization constants are used but with a deterministic center crop, which keeps the train and deploy distributions aligned. Because both exterior and wrist views share one encoder, they also share this preprocessing exactly.
Action Representation
Actions are expressed as seven-dimensional deltas covering three translation axes, three rotation axes and a scalar gripper command. Predicting deltas rather than absolute targets makes the policy less sensitive to calibration drift and lets it compose motions incrementally. Inactive action dimensions are masked out of the loss so that embodiments with partial actuation do not introduce spurious gradients.
Evaluation
Validation is performed on episodes that are held out from training, with metrics computed in the same normalized space used for optimization so that results are comparable across runs. The auxiliary future-state prediction is monitored as a diagnostic of whether the latent plan is becoming genuinely predictive rather than collapsing to a constant. Because inference is deterministic given a fixed seed and observation, reported behaviors are reproducible.
Reproducibility
Every run is made repeatable through explicit seeding of the data loader ordering, the augmentation transforms and the parameter initialization. Checkpoints are written atomically and include enough state to resume mid-epoch without divergence. The configuration embedded in the published artifact records the exact architectural choices, so loading the weights reconstructs the network without external information.
Adapting the Model
The modular structure invites adaptation. New vision backbones, language encoders or action parametrizations can be introduced by implementing the same token-level interface, and the number of plan tokens or thinking steps can be scaled to study the value of deliberation. New observation modalities are supported by adding a corresponding preprocessing branch that emits tokens of the shared width, after which the existing fusion and control stages operate unchanged.
- Downloads last month
- 23
