### Install villa-x Package

Source: https://context7.com/lucidrains/villa-x/llms.txt

Install the ViLLa-X library using pip.

```bash
pip install villa-x
```

--------------------------------

### Initialize and Use ViLLaX Model

Source: https://context7.com/lucidrains/villa-x/llms.txt

Shows how to initialize the LatentActionModel, ACTLatent, and ACTRobot components, compose them into the full ViLLaX model, and perform end-to-end inference with optional VLM and wrist image conditioning.

```python
import torch
from villa_x import LatentActionModel, ACTLatent, ACTRobot, ViLLaX

lam = LatentActionModel(
    dim=512,
    dim_proprio=17,
    dim_actions=20,
    patch_size=32,
    vit=dict(
        image_size=256,
        patch_size=32,
        num_classes=1000,
        dim=512,
        depth=6,
        heads=16,
        mlp_dim=2048,
        dropout=0.1,
        emb_dropout=0.1
    ),
)

act_latent = ACTLatent()
act_robot = ACTRobot(dim_proprio=11)

villa_x = ViLLaX(
    lam=lam,
    act_latent=act_latent,
    act_robot=act_robot
)

vlm_key_values = [
    (torch.randn(1, 4, 32, 64), torch.randn(1, 4, 32, 64)),
    (torch.randn(1, 4, 32, 64), torch.randn(1, 4, 32, 64)),
]

wrist_image = torch.randn(1, 3, 224, 224)

sampled_actions = villa_x(
    vlm_key_values=vlm_key_values,
    wrist_image=wrist_image
)
print(f"Generated actions shape: {sampled_actions.shape}")
```

--------------------------------

### Initialize and Use ACTLatent for Latent Action Generation

Source: https://context7.com/lucidrains/villa-x/llms.txt

Initialize ACTLatent for generating latent action sequences using flow matching, conditioned on VLM key/value pairs. Use for training by computing flow matching loss, or for inference by sampling latent actions.

```python
import torch
from villa_x import ACTLatent

# Initialize ACTLatent with default flow transformer settings
act_latent = ACTLatent(
    flow_dit=dict(
        dim_input=128,          # Latent action dimension
        dim_time=512,           # Time embedding dimension
        transformer=dict(
            dim=512,
            depth=6,
            heads=8,
            attn_dim_head=64,
        )
    )
)

# Optional: VLM key/value pairs from a vision-language model (e.g., PaLiGemma)
# List of (key, value) tuples from transformer layers
vlm_key_values = [
    (torch.randn(1, 4, 32, 64), torch.randn(1, 4, 32, 64)),  # Layer 1
    (torch.randn(1, 4, 32, 64), torch.randn(1, 4, 32, 64)),  # Layer 2
]

# Training: compute flow matching loss on action latents
action_latents = torch.randn(1, 32, 128)  # (batch, seq_len, latent_dim)
loss = act_latent(action_latents, vlm_key_values=vlm_key_values)
loss.backward()

# Inference: sample latent actions
# Requires specifying sequence length via shape parameter
sampled_latents = act_latent.sample(
    seq_len=32,
    vlm_key_values=vlm_key_values
)
print(f"Sampled latents shape: {sampled_latents.shape}")
```

--------------------------------

### Initialize and Use LatentActionModel

Source: https://context7.com/lucidrains/villa-x/llms.txt

Initialize the LatentActionModel with a ViT backbone for encoding video into latent action tokens. Use for training with video, proprioception, and actions, or for inference with video only.

```python
import torch
from villa_x import LatentActionModel
from vit_pytorch import ViT

# Initialize the Latent Action Model with a Vision Transformer backbone
lam = LatentActionModel(
    dim=512,                    # Model dimension
    dim_proprio=17,             # Proprioception dimension (joint states, etc.)
    dim_actions=20,             # Action dimension (motor commands)
    patch_size=32,              # Image patch size for reconstruction
    vit=ViT(
        image_size=256,
        patch_size=32,
        num_classes=1000,
        dim=512,
        depth=6,
        heads=16,
        mlp_dim=2048,
        dropout=0.1,
        emb_dropout=0.1
    ),
)

# Training: provide video, proprioception, and actions
# video shape: (batch, channels, time, height, width)
video = torch.randn(2, 3, 8, 256, 256)
proprio = torch.randn(2, 8, 17)   # (batch, time, proprio_dim)
actions = torch.randn(2, 8, 20)   # (batch, time, action_dim)

# Forward pass returns loss and breakdown (ar_fdm_loss, ar_proprio_loss, ar_action_loss)
loss, (fdm_loss, proprio_loss, action_loss) = lam(video, proprio, actions)
loss.backward()

# Inference: extract latent action tokens from video only
latent_action_tokens = lam(video)
# Returns shape: (batch, time-1, dim) = (2, 7, 512)
print(f"Latent tokens shape: {latent_action_tokens.shape}")
```

--------------------------------

### Training and Inference with ACTRobot

Source: https://context7.com/lucidrains/villa-x/llms.txt

Demonstrates how to compute the flow matching loss during training and sample robot actions during inference using the ACTRobot component, with optional VLM conditioning.

```python
vlm_key_values = [
    (torch.randn(1, 4, 32, 64), torch.randn(1, 4, 32, 64)),
    (torch.randn(1, 4, 32, 64), torch.randn(1, 4, 32, 64)),
]

actions = torch.randn(1, 128, 20)  # (batch, seq_len, action_dim)
loss = act_robot(
    actions,
    action_latents,
    proprio=proprio,
    wrist_image=wrist_image,
    vlm_key_values=vlm_key_values
)
loss.backward()

sampled_actions = act_robot.sample(
    action_latents,
    wrist_image=wrist_image,
    vlm_key_values=vlm_key_values
)
print(f"Sampled actions shape: {sampled_actions.shape}")
```

--------------------------------

### Configure and Use FlowTransformerWrapper

Source: https://context7.com/lucidrains/villa-x/llms.txt

Demonstrates creating a custom FlowTransformerWrapper with detailed configuration for diffusion transformer architecture, including VLM conditioning and cross-attention. It also shows wrapping it with RectifiedFlow for training and sampling.

```python
import torch
from villa_x.villa_x import FlowTransformerWrapper
from rectified_flow_pytorch import RectifiedFlow

flow_dit = FlowTransformerWrapper(
    dim_input=128,
    dim_time=512,
    transformer=dict(
        dim=512,
        depth=6,
        heads=8,
        attn_dim_head=64,
    ),
    cross_attend=True,
    cross_attn_dim_context=128,
    dropout_vlm_key_values=0.5,
)

flow_wrapper = RectifiedFlow(flow_dit)

action_latents = torch.randn(2, 32, 128)
context = torch.randn(2, 16, 128)

loss = flow_wrapper(action_latents, context=context)
loss.backward()

sampled = flow_wrapper.sample(seq_len=32, context=context)
print(f"Sampled shape: {sampled.shape}")
```

--------------------------------

### Initialize and Use ACTRobot for Robot Action Decoding

Source: https://context7.com/lucidrains/villa-x/llms.txt

Initialize ACTRobot to decode latent actions into robot motor commands, optionally conditioned on wrist camera images and proprioceptive state. This component uses flow matching for precise manipulation control.

```python
import torch
from villa_x import ACTRobot

# Initialize ACTRobot with proprioception support
act_robot = ACTRobot(
    dim_proprio=11,             # Proprioception dimension
    dim_action_latent=128,      # Latent action dimension (must match ACTLatent)
    flow_dit=dict(
        dim_input=20,           # Output action dimension
        transformer=dict(
            dim=512,
            depth=6,
            heads=8,
        )
    )
)

# Prepare conditioning inputs
action_latents = torch.randn(1, 32, 128)    # From ACTLatent
proprio = torch.randn(1, 11)                 # Current proprioceptive state
wrist_image = torch.randn(1, 3, 224, 224)   # Wrist camera image (optional)


```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.