UniAction: Unified Action Conditioning for World Models

Paper: UniAction: Unified Action Conditioning for World Models
Authors: Bowen Xue, Zizhang Li, Hong-Xing Yu, Jiajun Wu
Affiliation: Stanford University
Paper: https://arxiv.org/abs/2610.32977
Project page: https://uniaction.github.io/
Code status: The official repository is public. This analysis was checked against commit b0d09151335f0682af4d6f5858ae188ed1c64981.

One-Sentence Summary

UniAction represents forces, camera commands, hand poses, and robot poses as sparse action elements with separate semantic identities and spatiotemporal coordinates, allowing one video world model to receive all four control forms through a shared native-attention pathway.

Why Action Conditioning Is Fragmented

A video world model predicts future observations from a current frame, a task description, and an action. Large video generators are a natural starting point because pretraining exposes them to object appearance, scene layout, contact, deformation, and motion. The hard part is not only learning video dynamics. It is building an action interface that can accept radically different control signals without designing a new model for every source.

A force is a vector applied at a point and time. A camera action is a discrete key press such as forward translation or yaw rotation. A hand pose is a dense sequence of per-frame joint coordinates. A robot pose can contain a different number of joints and links depending on the embodiment. Some actions are short events, while others persist across many frames. Some have an obvious image location, while others affect the whole viewpoint.

Most existing action-conditioned video models specialize in one of these forms. The encoder, conditioning layer, preprocessing pipeline, and training set are tailored to a single action family. This is practical in isolation, but the resulting collection of models wastes a scarce resource: action-annotated video. Each dataset can improve only one specialist, even though all four domains contain evidence about materials, occlusion, gravity, object permanence, and scene dynamics.

Visual conditioning offers one generic route. An action is rendered into an image or video control signal, encoded by a VAE, and injected through a conditioning branch. This gives strong spatial correspondence when a natural rendering exists, such as a skeleton or a highlighted key. It is less natural for abstract interventions. A force has no canonical appearance, and rendering it as an arrow introduces design choices that are not part of the physical action.

Token conditioning is another generic route. A separate encoder maps each action form into tokens, and cross-attention injects those tokens into the video model. This avoids explicit rendering, but the relation between an action token and the frame or pixel where it takes effect remains implicit in learned features. Pose control requires precise localization, so this indirect relation can become the bottleneck.

The paper starts from a compact observation: regardless of the raw form, a world model needs to know two things about an action. It needs the action's meaning, and it needs the time and place where that meaning affects the world. Once these two pieces are represented separately, heterogeneous controls can share one interface.

The Unified Representation

UniAction keeps the usual conditional video-generation objective. The model learns the distribution of a future video conditioned on the first frame, text, and action.

$$p_\theta(V\mid I_0,T,A)$$

The central change is how the action $A$ is described. Every action becomes a sparse set of $N$ elements. Element $i$ contains a discrete semantic identity $s_i$ and a continuous coordinate $\mathbf p_i$.

$$A=\{(s_i,\mathbf p_i)\}_{i=1}^{N},\qquad \mathbf p_i=(t_i,y_i,x_i)$$

Here $t_i$ is the time at which the element takes effect, while $y_i$ and $x_i$ locate it in the video frame. A force can use two elements. A dense robot trajectory can use many. The raw numbers may differ, but every action is converted into the same structural contract.

flowchart LR
  F[Force: two vector endpoints] --> S[Sparse action elements]
  C[Camera: eight command identities] --> S
  H[Hand: visible per-frame joints] --> S
  R[Robot: visible embodiment joints] --> S
  S --> E[Shared semantic codebook]
  S --> P[Time and image coordinate]
  E --> T[Concatenate with video tokens]
  P --> T
  T --> ROPE[Map into native 3D RoPE]
  ROPE --> CAL[Calibrate action attention mass]
  CAL --> M[One shared video world model]
    

Figure 1: The teaser contrasts four specialist models with one shared model. The red bars show that the jointly trained UniAction model improves action following for all four action forms.

Figure 1: One world model for heterogeneous actions versus separate action-specific experts.

Semantic identity: a shared codebook

The discrete part of an action indexes a shared semantic vocabulary. A learned codebook $E$ maps each identity to a token whose dimension matches the video transformer.

$$E\in\mathbb R^{|\mathcal S|\times d},\qquad \mathbf a_i=E[s_i]\in\mathbb R^d$$

This factorization is important. The codebook says what an element means, while its coordinate says where and when the element acts. A hand joint at two positions can reuse the same semantic token. A force start and force endpoint use different identities even when their coordinates are close or identical. The model does not need to learn a separate raw-action encoder for each family.

The vocabulary can contain semantic roles across families: force start, force endpoint, camera directions, hand joints, and robot joints. New identities can be added without replacing the video backbone. The paper therefore treats extensibility as a property of the representation rather than as a new network head for every action type.

Spatiotemporal grounding: reuse the backbone's geometry

Meaning alone is insufficient. The same joint movement at a different pixel should produce a different visual effect, and the same force applied later should trigger motion later. UniAction maps action coordinates into the same positional system that the video transformer already learned during pretraining.

Let $\Phi_{\mathcal B}$ denote the backbone's positional mechanism. An action token and a video token are encoded as follows:

$$\widetilde{\mathbf a}_i=\Phi_{\mathcal B}(\mathbf a_i,\mathbf p_i),\qquad \widetilde{\mathbf v}_j=\Phi_{\mathcal B}(\mathbf v_j,\mathbf p_j^v)$$

A cross-attention baseline can also encode time and position inside its action encoder. The difference is where that information lives. UniAction places action tokens in the pretrained model's own spatiotemporal geometry. Coordinates are first converted by the backbone's VAE compression and tokenization:

$$\mathbf p_i^{\mathrm{lat}}=g_{\mathcal B}(\mathbf p_i)$$

The action tokens are then appended to the latent video sequence and processed by the backbone's native self-attention.

$$[\mathbf v_1,\ldots,\mathbf v_M,\mathbf a_1,\ldots,\mathbf a_N]$$

This gives UniAction a stronger inductive bias than merely supplying learned position vectors. When the video backbone has already learned how temporal and spatial rotary phases relate, action coordinates can reuse that structure. The paper reports that replacing this geometry with sinusoidal or learned absolute position embeddings causes a large loss in precise pose control.

Figure 2: The method converts four action forms into semantic identities and coordinates, grounds both action and video tokens in the backbone geometry, and uses cardinality-calibrated attention.

Figure 2: The UniAction conditioning pipeline.

How the four action forms are instantiated

For hand pose, each semantic identity is a joint index. Every visible joint at every frame yields one action element at the joint's image location. Robot pose follows the same principle, with identities defined by the joints of the specific robot. The model does not force two embodiments into one shared skeleton before training.

Camera control uses eight discrete keys: forward, backward, left, right, pitch up, pitch down, yaw left, and yaw right. The command has no intrinsic image location, so each direction receives a fixed canonical location at the corresponding side of the frame. Translation and rotation with the same direction share that location but retain different semantic identities.

Force uses two tokens because one point cannot encode both application point and vector direction. One token marks the tail and the other marks the head. Their spatial offset is the force vector, while their semantic roles distinguish the two endpoints. The official implementation rejects a force sample unless it contains exactly one start and one endpoint token.

Cardinality-calibrated attention

Once all tokens share one self-attention sequence, representation granularity creates a new bias. A robot pose may contribute hundreds of joint-time tokens, while a force contributes two. The attention matrix naturally gives the larger set more total mass, even if each token has similar compatibility with the video query.

For an action with $N$ tokens, the total unnormalized attention mass is:

$$Z_a=\sum_{i=1}^{N}\exp(\ell_i^a)=N\cdot\operatorname{mean}_i[\exp(\ell_i^a)]$$

The factor $N$ depends only on representation cardinality. To remove it, UniAction adds the same offset to every action logit:

$$\hat\ell_i^a=\ell_i^a+\log\frac{N_{\mathrm{ref}}}{N}$$

Because the offset is constant across the action keys, their relative ordering does not change. The total action mass now scales with the fixed reference cardinality $N_{\mathrm{ref}}$, which is 784 in the main experiments. The calibrated softmax and output become:

$$[A_v\mid A_a]=\operatorname{softmax}([L_v\mid L_a+b]),\qquad b=\log\frac{N_{\mathrm{ref}}}{N_a}$$

$$O_v=A_vV_v+A_aV_a$$

This is a small change with a structural purpose. It lets the shared model compare sparse and dense actions on a common attention scale. Without it, dense pose tokens dominate; with it, a two-token force can still steer the model strongly enough.

Experimental Design

The main training set contains 10,000 clips from each of GoalForce, SolarWM, EgoDex, and DROID, for 40,000 clips total. Every clip is sampled to 81 frames and resized or cropped to 832 by 480. The GoalForce portion includes 5,000 ball, 2,500 domino, and 2,500 flower samples.

The generative base is the high-noise expert of Wan2.2-I2V-A14B. LoRA is used for adaptation with rank and alpha equal to 128. During action adaptation, only the action branch is updated while the video branch remains frozen. The main UniAction setup has about 404 million trainable parameters and is trained for two passes over the data.

Evaluation fixes the initial frame and text prompt, then varies only the action. The authors construct 410 counterfactual generations from 107 source samples. A ball force can be rotated by 45 or 90 degrees, a domino force can be reversed, a pose trajectory can be flipped or frozen, and a camera command can be changed. This design tests whether the generated response changes with the intervention instead of only matching a familiar scene.

DIS optical flow measures force-induced object motion. CoTracker3 tracks hand keypoints and robot control points. ViPE estimates camera trajectories. The force direction score uses the cosine similarity between predicted and target directions:

$$s_{\mathrm F}=C\left(\frac{\hat{\mathbf f}^{\top}\mathbf f}{\|\hat{\mathbf f}\|\,\|\mathbf f\|}\right)$$

Here $C(x)$ clips the score to the unit interval. A sample with no valid motion response receives zero. The paper also reports AlexNet LPIPS against original-condition ground truth and a human study with 60 cases and 100 participants.

Main Comparison

UniAction is compared with VACE-style visual conditioning and a cross-attention token baseline. All methods start from the same pretrained model and receive matched training data and optimization. The comparison therefore isolates how actions are represented and injected.

Action	Method	Action Following↑	LPIPS↓	Human Following↑	Human Quality↑
Force	VACE	0.6046	0.1178	2.47	3.40
Force	Cross-Attn	0.1618	0.1341	1.96	3.13
Force	UniAction	0.7630	0.1163	4.37	4.00
Camera	VACE	0.4997	0.5100	3.30	3.55
Camera	Cross-Attn	0.5306	0.5124	3.13	3.35
Camera	UniAction	0.6371	0.4863	4.02	3.68
Hand	VACE	0.2938	0.1122	3.30	3.43
Hand	Cross-Attn	0.0678	0.1695	2.34	3.12
Hand	UniAction	0.4119	0.0910	3.36	3.42
Robot	VACE	0.5245	0.1389	3.60	3.73
Robot	Cross-Attn	0.0859	0.1859	2.07	3.44
Robot	UniAction	0.6521	0.1164	4.32	3.99

The strongest gaps appear on force, hand pose, and robot pose. Cross-attention is especially weak on hand and robot localization. This fits the paper's argument that a learned encoder has to rediscover spatial correspondence from limited action data. VACE remains competitive when a natural visualization exists, but it cannot give force or keyboard commands a canonical visual form without a handcrafted renderer.

Figure 3: Across forces, camera commands, hand poses, and robot poses, UniAction follows the overlaid action more reliably than the two generic baselines.

Figure 3: Qualitative comparison of alternative conditioning mechanisms.

Joint Training versus Four Experts

The second experiment asks whether heterogeneous data help one another. The authors train a separate expert for each action form. Every expert has the same architecture as UniAction and sees the same samples as the joint model. The four experts together use about four times the parameters and hard-route each sample by action type.

Action	Expert Following↑	Joint UniAction↑	Expert LPIPS↓	Joint UniAction↓
Force	0.5436	0.7630	0.1295	0.1163
Camera	0.5182	0.6371	0.4780	0.4863
Hand	0.3701	0.4119	0.0934	0.0910
Robot	0.6310	0.6521	0.1286	0.1164

UniAction wins action following on all four tasks. Its LPIPS is lower on three tasks and only 0.0083 higher on camera control. This result is more informative than a comparison against an unrelated baseline because the joint model and experts use matched data and compute. The gain is consistent with knowledge transfer: force data can teach the model about materials and contact, camera data can teach viewpoint and scene layout, and manipulation data can teach articulation and object interaction.

Figure 4: The second row in each pair is the jointly trained model. It produces more faithful dynamics than the action-specific expert while preserving plausible appearance.

Figure 4: Joint training compared with four action-specific experts.

Ablations and Temporal Generalization

The ablations replace the native 3D RoPE with sinusoidal or learned absolute position embeddings, move grounded action tokens through cross-attention, or remove cardinality calibration. All variants use Wan2.2-TI2V-5B under the same protocol.

Variant	Force	Camera	Hand	Robot
Self-attention to Cross-attention	0.2975	0.5050	0.0533	0.0733
RoPE to Sinusoidal APE	0.2084	0.5278	0.0603	0.0755
RoPE to Learned APE	0.1485	0.5206	0.0656	0.0693
Without Cardinality Calibration	0.5186	0.4966	0.4303	0.6995
UniAction	0.7190	0.5842	0.4266	0.7004

The geometry variants nearly fail on hand and robot pose. Merely providing coordinates is not enough if those coordinates live in a different positional space from the video tokens. The pretrained rotary geometry gives the action a reusable temporal and spatial structure, while learned absolute embeddings must infer that structure from the much smaller action dataset.

The cardinality result has a useful tradeoff. Without calibration, dense hand and robot actions score slightly higher, but sparse force and camera actions collapse. This is the predicted bias: many pose tokens dominate attention, while few force tokens are suppressed. Calibration sacrifices a small amount of dense-action mass to make all action families controllable within one model.

The force temporal experiment makes the geometry argument concrete. During training, all forces occur at frame zero. At test time, the authors move the force to frame 8 or frame 24. UniAction starts the induced object motion at the newly specified time. A learned absolute position embedding barely shifts because the unseen time bins received no supervision.

Figure 5: Both variants are trained with forces only at frame zero. UniAction delays the response when the test force is moved to frame 8 or 24, while learned APE does not transfer to unseen time slots.

Figure 5: Temporal generalization to forces at unseen frame indices.

Efficiency tradeoffs

UniAction has far fewer trainable parameters than VACE or the cross-attention baseline: about 404 million, compared with 3.1 billion and 640 million. Its added latency is also the lowest for all four action types. The token-count effect remains visible, however: hand conditioning adds 72.9 TFLOPs because many joints enter self-attention, while a force adds only 0.1.

Method	Parameters	Added TFLOPs	Added Latency
VACE	3.1B	351.2 for all four	about 21%
Cross-Attn	640M	33.5 to 36.0	8.6% to 9.0%
UniAction force	404M	0.1	3.6%
UniAction camera	404M	5.9	3.0%
UniAction hand	404M	72.9	6.6%
UniAction robot	404M	37.2	4.9%

Figure 6: An additional force comparison. UniAction preserves the initial domino push direction and the resulting chain reaction more consistently.

Figure 6: Additional qualitative force conditioning result.

What the Code Reveals

The public implementation exposes the operational constraints behind the paper. The action configuration fixes the number of frames, image size, semantic vocabulary sizes, force identities, and reference cardinality. These values are not only appendix details; they define tensor shapes and validated input paths.

# src/uniaction/model/action.py:22-37
@dataclass(frozen=True)
class UnifiedActionConfig:
    l_max: int = 32
    num_frames: int = 81
    dit_dim: int = 5120
    video_height: int = 480
    video_width: int = 832
    hand_vocab_size: int = 16
    camera_vocab_size: int = 16
    robot_vocab_size: int = 12
    force_start_token_id: int = 0
    force_endpoint_token_id: int = 3
    action_reference_count: float = 784.0

The force encoder enforces the two-point representation at runtime. It requires a start token and an endpoint token, rejects any other count, and checks that both tokens use time zero for their positional encoding. The start and endpoint then receive separate learned semantic tokens.

# src/uniaction/model/action.py:191-209
if count != 2:
    raise ValueError(f"force requires exactly two tokens, got {count}")
start = ids == self.force_start_token_id
endpoint = ids == self.force_endpoint_token_id
if not bool(((start.sum(dim=1) == 1) & (endpoint.sum(dim=1) == 1)).all()):
    raise ValueError("force requires one start token and one endpoint token")
if not bool((time_ids.float() == 0).all()):
    raise ValueError("force start and endpoint RoPE must both use time zero")
tokens = torch.where(
    start.unsqueeze(-1),
    start_token.expand(batch, count, -1),
    endpoint_token.expand(batch, count, -1),
)
return self.force_token_norm(tokens)

The grounding path is implemented most directly in `_action_freqs()`. The method infers a three-dimensional latent grid, maps normalized UV coordinates to height and width positions, maps action time to the latent frame grid, and creates rotary frequencies along all three axes. This is the concrete code form of the paper's $g_{\mathcal B}$ and $\Phi_{\mathcal B}$ notation.

# src/uniaction/model/action.py:452-476
t_pos = time_ids * (float(f - 1) / float(adapter.num_frames - 1))
v_norm = uv[:, 1].clamp(0.0, 1.0)
u_norm = uv[:, 0].clamp(0.0, 1.0)
h_pos = (v_norm * float(h) - 0.5)
w_pos = (u_norm * float(w) - 0.5)
state.freqs = torch.cat(
    [
        wan_rope_at_positions(t_pos, dim_t),
        wan_rope_at_positions(h_pos, dim_h),
        wan_rope_at_positions(w_pos, dim_w),
    ],
    dim=-1,
)

Cardinality calibration is a single log-weight calculation. The reference count is stored as a non-trainable log scalar, and the implementation subtracts the logarithm of the current action-token count. No extra learned action classifier is needed.

# src/uniaction/model/action.py:510-520
def action_key_log_weight(
    *,
    video_token_count: int,
    action_token_count: int,
    action_log_reference: torch.Tensor,
) -> torch.Tensor:
    if video_token_count <= 0 or action_token_count <= 0:
        raise ValueError("video and action token counts must be positive")
    if action_log_reference.numel() != 1:
        raise ValueError("action_log_reference must be scalar")
    return action_log_reference.reshape(()) - math.log(float(action_token_count))

The training code also separates the stages cleanly. Stage 1 trains and merges shared video LoRA weights. Stage 2 builds the action adapter, attaches it to every Wan DiT expert, injects action-branch LoRA, and freezes the reference-count parameter. The trainable schema is validated by tensor count and dtype so an accidentally unfrozen backbone component cannot pass silently.

Limitations

The authors explicitly motivate the work with the scarcity of action-annotated data and its concentration in a narrow range of scenes and interactions. That concentration remains a limitation after unification. A shared model can pool four datasets, but it cannot invent experiences absent from GoalForce, SolarWM, EgoDex, or DROID.

The evaluation covers four hand-designed action forms and a limited set of semantic identities. It does not establish open-world action vocabularies, arbitrary robot embodiments, widely varying control frequencies, or long-horizon closed-loop planning. Each new form may still require a meaningful mapping from raw signals to semantic identities and grounded coordinates.

The primary metric is generated-video action following, not physical robot task success. The study does not place the generated futures in a control loop and measure whether a policy improves manipulation success. This distinction matters because plausible video can still violate contact constraints or produce dynamics that are not executable on hardware.

Automatic scores depend on off-the-shelf motion estimators. Optical flow can fail under occlusion, CoTracker3 can lose keypoints, and ViPE can confuse camera motion with scene motion. The paper reduces camera drift by subtracting a no-action reference, but estimator error still remains. The human study helps, yet it rates visual following rather than physical validity.

Finally, the approach depends on a large video backbone and action-specific preprocessing. The conditioning overhead is smaller than the baselines, but the A14B base remains expensive. Hand actions add many tokens and substantial computation, so scaling to denser skeletons or longer videos will require further token compression or hierarchical attention.

Conclusion

UniAction's main contribution is a representation-level factorization: semantic identity explains what an action means, and a continuous coordinate explains where and when it acts. By grounding both action and video tokens in the pretrained 3D RoPE geometry, the method reuses learned dynamics instead of relearning correspondence from a small action dataset. Cardinality calibration then prevents token count from becoming an accidental control-strength prior.

The strongest evidence is the joint-versus-expert result. One model beats four matched specialists across all action forms, indicating that heterogeneous action data can share useful world knowledge. The result does not yet prove general robot intelligence, but it provides a credible interface for adding new intervention types to one continually expanding video world model.

Golden quote: Unifying actions does not mean erasing their raw forms. It means preserving what each action means, placing it in the world's time and space, and giving every action the same way to speak to the model.