reference/generated/torchrl.modules.vla.VLAWrapperBase

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

VLAWrapperBase

class torchrl.modules.vla.VLAWrapperBase(*args, **kwargs)

Base class for TensorDict-native Vision-Language-Action policies.

A VLA policy maps images, optional proprioceptive state, and a language instruction to either a continuous action chunk or discrete action tokens. Outputs are stored in a structured VLAAction container under "vla_action" by default. Its fields are ordinary nested TensorDict keys, e.g. ("vla_action", "chunk") for continuous chunks.

Keyword Arguments:

• action_dim (int) – The dimensionality of a single action.

• chunk_size (int) – The action-chunk horizon.

• action_head (str) – "continuous" or "tokens".

• input_mode (str) – "canonical" reads raw VLA keys. "preprocessed" reads a VLAObservation or TensorDictBase from observation_key.

• output_mode (str, optional) – "chunk", "tokens" or "both". Defaults to "chunk" for continuous heads and "tokens" for token heads.

• return_vla_action_container (bool) – whether to write the structured VLAAction object at the VLA action root key. When False, only its plain TensorDict fields are written. Defaults to True.

• vocab_size (int, optional) – Number of action-token bins, required for token heads.

• action_tokenizer (ActionTokenizerBase, optional) – Token/chunk codec used when output_mode asks for both representations.

• return_log_probs (bool, optional) – Whether token forward writes log-probabilities. Defaults to True for token heads.

• return_logits (bool) – Whether token forward writes action_logits.

• logits_only (bool) – Whether token forward returns logits without sampling actions by default. A per-call logits_only=True argument also enables this path.

• log_probs_mode (str) – "sequence" returns one summed log-probability per sample; "token" returns per-token log-probabilities.

• use_state (bool) – Whether canonical mode reads the state key.

• default_interaction_type (InteractionType) – Token readout when no exploration context is active.

• mode (str, optional) – Backward-compatible alias mapping "sample" to InteractionType.RANDOM and "greedy" to deterministic.

• inplace (bool | "empty" | None) – Output TensorDict behavior. True updates the input, False returns a new output TensorDict, and "empty" returns an empty TensorDict populated with outputs.

• num_samples (int, optional) – Number of token samples to draw per input.

Examples

>>> import torch
>>> from tensordict import NonTensorStack, TensorDict
>>> from torchrl.modules.vla import TinyVLA
>>> policy = TinyVLA(action_dim=7, chunk_size=4)
>>> td = TensorDict(
...     {
...         "observation": {
...             "image": torch.zeros(2, 3, 16, 16, dtype=torch.uint8),
...             "state": torch.zeros(2, 5),
...         },
...         "language_instruction": NonTensorStack("pick", "place"),
...     },
...     batch_size=[2],
... )
>>> out = policy(td)
>>> out["vla_action"].chunk.shape
torch.Size([2, 4, 7])
>>> out["vla_action", "chunk"].shape
torch.Size([2, 4, 7])

forward(tensordict: TensorDictBase, *, tensordict_out: TensorDictBase | None = None, logits_only: bool = False, **kwargs) → TensorDictBase

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_dist(tensordict: TensorDictBase, *, tensordict_out: TensorDictBase | None = None, logits_key: NestedKey | None = None, mask_key: NestedKey | None = None, **kwargs) → Distribution

Return the action-token distribution for loss-time recomputation.

get_new_version(**kwargs) → VLAWrapperBase

Return a shallow wrapper copy with altered runtime parameters.

log_prob(tensordict: TensorDictBase, *, action_key: NestedKey | None = None, log_probs_key: NestedKey | None = None, **kwargs) → TensorDictBase

Recompute and write token log-probabilities for stored actions.

set_keys(**kwargs) → VLAWrapperBase

Set the tensordict key names used by the policy.