JAST-DiT: Joint Action-State-Text Diffusion

• The paper introduces a unified diffusion model that integrates action, state, and text streams through coordinated tokenization and cross-modal attention.
• It employs a multi-stage training pipeline, combining supervised pretraining with reinforcement learning and hybrid rewards to enhance physical realism and semantic alignment.
• Empirical evaluations demonstrate state-of-the-art performance in language-conditioned motion control, achieving improved fidelity and inference speed across benchmarks.

A Joint Action-State-Text Diffusion Transformer (JAST-DiT) is a class of generative neural architectures that learns the conditional distribution over multi-modal trajectories consisting of action, physical state, and text or language context, using a unified diffusion-based modeling framework. JAST-DiT models are structured to process these distinct modalities as coordinated streams or unified token sequences, enabling direct semantic and dynamical interaction, and are deployed for tasks such as physically-plausible motion control from language instructions, multi-modal sequence generation, and vision-language-action reasoning. Representative implementations include the JAST-DiT core of SCRIPT for language-conditioned humanoid control (Zhang et al., 21 May 2026), the three-way variant embedded in PackDiT for bidirectional motion–text generation (Jiang et al., 27 Jan 2025), and extensions to vision-language-action joint reasoning as in Unified Diffusion VLA (Chen et al., 3 Nov 2025).

1. Architecture and Tokenization

JAST-DiT models are characterized by either multi-stream or unified sequence tokenization strategies, depending on the application domain and the modality interaction required. In SCRIPT (Zhang et al., 21 May 2026), JAST-DiT maintains three explicit streams:

• Action tokens ($\mathbf{a}_{t+1:t+H}$): future control targets (e.g., joint actuation vectors),
• State tokens ($\mathbf{s}_{t+1:t+H}$): proprioceptive or world state observations (e.g., joint positions, velocities),
• Text tokens ($\mathbf{c}_{\text{txt}}, \mathbf{c}_{\text{pool}}$): CLIP-encoded penultimate-layer features and pooled text representations.

Each stream undergoes linear projection to a shared model dimension, is augmented with stream-specific positional encodings, and receives global conditioning ($$\mathbf{e} = \text{Cond}(\tau, \mathbf{c}_{\text{pool}})$$) via AdaLN-Zero before self-attention.

PackDiT's extension formulates three parallel diffusion transformer (DiT) backbones—one per modality (action, state, text)—with cross-modal "mutual blocks" applying attention between all pairs (A↔S, S↔T, A↔T) (Jiang et al., 27 Jan 2025). Unified Diffusion VLA (Chen et al., 3 Nov 2025) instead casts all modalities (text instructions, current/future vision, actions) into a single discrete token sequence subject to a hybrid attention mask.

2. Joint Attention and Cross-modal Interaction

JAST-DiT enables semantic and dynamical entanglement of its modalities through specialized attention mechanisms. In SCRIPT (Zhang et al., 21 May 2026), within each Transformer block:

• Queries, keys, and values are computed independently for actions, states, and text.
• These are then concatenated and fed through multi-head joint self-attention:

$$[\bar{\mathbf{A}}^{(\ell)};\;\bar{\mathbf{S}}^{(\ell)};\;\bar{\mathbf{C}}^{(\ell)}] = \mathrm{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V})$$

• No causal mask is applied to action/state streams, while the standard padding mask is used for text.

After joint attention, nonlinear history-conditioned cross-attention injects sampled context from dense recent and sparse distant frames, formulated as a two-stage cross-attentive retrieval (long-range then dense recent).

PackDiT implements interleaved mutual prompting blocks wherein each modality's DiT attends to another, aligning representations and supporting bi-directional (e.g., text→motion and motion→text) generative capabilities.

Unified Diffusion VLA's hybrid attention scheme prescribes modality-specific attention flows—causal for text, bidirectional for intra-vision/action, and causal from instructions/observations to outputs—enforced via a bespoke attention mask (Chen et al., 3 Nov 2025).

3. Diffusion and Denoising Objectives

JAST-DiT leverages model-based diffusion processes for joint generative modeling:

• SCRIPT (Zhang et al., 21 May 2026) uses flow-matching diffusion on the concatenated future (action+state) vectors; noise is injected to $\mathbf{x}_0 \sim \mathcal{N}(0, I)$, and the model learns a vector field $$v_\theta(\mathbf{x}_\tau, \tau, \mathcal{H}_t, \mathbf{c})$$ to match the target data under mean-squared error:

$$\mathcal{L}_{\text{FM}} = \mathbb{E}_{\tau, \mathbf{x}_0, \mathbf{x}_1, \mathcal{H}_t, \mathbf{c}} \| v_\theta(\mathbf{x}_\tau, \tau, \mathcal{H}_t, \mathbf{c}) - (\mathbf{x}_1 - \mathbf{x}_0) \|_2^2$$

• PackDiT (Jiang et al., 27 Jan 2025) employs standard DDPM-style Gaussian diffusion per modality, with losses summed across action (A), state (S), and text (T) components.
• Unified Diffusion VLA (Chen et al., 3 Nov 2025) adopts a joint discrete denoising diffusion process (JD3P). Target (future frames + actions) tokens are randomly masked per diffusion step; model training minimizes cross-entropy over masked tokens, and parallel denoising is performed synchronously at inference.

4. Training Pipelines and Optimization

JAST-DiT-based models use multi-stage training to leverage both offline data and online reinforcement learning:

• Stage I: Supervised Imitation / Pretraining Models are pre-trained to minimize diffusion losses (e.g., $\mathcal{L}_{\text{FM}}$ for SCRIPT) on curated action-state-text trajectories. In SCRIPT, classifier-free guidance is instated by randomly dropping text streams.
• Stage II: Reinforcement Learning with Hybrid Rewards (RLHR) Following pretraining, parameters are refined with RL using hybrid reward signals, including dense physical alignment and sparse text-alignment rewards (computed via frozen state–text contrastive models). The stochastic flow sampler is treated as a policy, optimized via PPO, value loss, entropy regularization, and a behavior cloning anchor ($\mathcal{L}_{\text{BC}}$) (Zhang et al., 21 May 2026).

Unified Diffusion VLA (Chen et al., 3 Nov 2025) uses a two-phase recipe: initial post-training on future image prediction, then full joint diffusion with actions.

5. Inference and Decoding Strategies

Inference with JAST-DiT varies across implementations:

• SCRIPT generates action/state rollouts conditioned on history buffers and language, operating autoregressively but stabilized by nonlinear history-conditioned context windows.
• PackDiT supports synthesis and prediction tasks in multiple directions (text↔motion, unconditional, in-betweening); sampling proceeds with cross-modal attention enabled as appropriate for each direction (Jiang et al., 27 Jan 2025).
• Unified Diffusion VLA attains high throughput through parallel synchronous denoising, confidence-based unmasking, and KV-caching for context blocks, yielding up to 4× speedup over autoregressive decoders (Chen et al., 3 Nov 2025).

6. Empirical Performance and Ablations

Empirical evaluation across datasets such as HumanML3D (motion-text), MotionMillion (physics-based humanoid control), and embodied robot benchmarks (CALVIN, LIBERO, SimplerEnv) demonstrates that JAST-DiT architectures set new state-of-the-art results in instruction fidelity, motion quality, and closed-loop physical realism (Zhang et al., 21 May 2026, Jiang et al., 27 Jan 2025, Chen et al., 3 Nov 2025). Key findings include:

• All three streams—action, state, text—are essential for semantic and physical control.
• Nonlinear history context improves stability and effective hindsight usage.
• Hybrid rewards in RLHR prevent degenerate behaviors and enforce both semantic alignment and physical feasibility.
• Scaling up model dimension and depth yields monotonic improvements in text–motion alignment (R-Precision), motion quality (FID/MM-Dist), and diversity; e.g., SCRIPT-Huge achieves R-Precision Top-1 of 0.464 and FID 0.645 on MotionMillion (Zhang et al., 21 May 2026).
• Unified Diffusion VLA exhibits rapid (219 tokens/sec) and effective inference, outperforming prior autoregressive and partial unification methods (Chen et al., 3 Nov 2025).

7. Generalization and Extensions

JAST-DiT principles have been expanded to multi-directional conditional generation, as in PackDiT's variant that supports state-prediction, action supervision, and arbitrary subset conditioning at sampling—enabling flexible combinations of action, state, and text conditioning (Jiang et al., 27 Jan 2025). Unified Diffusion VLA generalizes this through inclusion of vision, permitting simultaneous understanding and generation tasks (read: text, see: images, act: future and action) within a single diffusion process and transformer backbone (Chen et al., 3 Nov 2025).

Significant ablation studies confirm the necessity of joint tokenization, attention hybrids, and multi-stage training for robust, scalable, and physically grounded multi-modal generative control. A plausible implication is that extending JAST-DiT design to even richer modalities (e.g., audio, tactile) or embedding further contextual knowledge may further enhance reasoning and generalization capabilities in embodied agents.