observation_tokenizer

observation_tokenizer

Observation tokenizer for creating unified prompts from multiple observation keys.

This tokenizer: 1. Takes multiple observation keys (language, proprio, gripper, etc.) 2. Optionally bins continuous data into quantiles 3. Creates a unified prompt string (e.g., "Task: grasp needle, State in robot frame: 127 143 89") 4. Tokenizes the prompt using a language model tokenizer 5. Returns token IDs and padding masks

ObservationTokenizer

ObservationTokenizer(tokenizer_model, observation_keys, bin_continuous_data=True, num_bins=256, max_token_len=256, device=None, raw_text=False, prompt_template=None, padding_strategy=value, trust_remote_code=False)

Tokenizes multiple observation keys into a unified prompt.

Creates prompts like: "Task: grasp needle, State in robot frame: 127 143 89, State in camera frame: 88 201 43"

All specified observation keys are combined into a single string, then tokenized.

Initialize observation tokenizer.

Parameters:

Name	Type	Description	Default
tokenizer_model	str	HuggingFace model name (e.g., "google/gemma-2b")	required
observation_keys	list[str]	List of observation keys to include in prompt (order preserved)	required
bin_continuous_data	bool	Whether to bin continuous data into quantiles	True
num_bins	int	Number of bins for quantile-based discretization	256
max_token_len	int	Maximum token sequence length	256
device	device | None	Target device for tensors	None
raw_text	bool	If True, pass language text through with only a trailing newline appended (no Task: prefix, no lowercasing). Use for VLM policies (SmolVLA, Pi0) that expect unformatted text.	False
prompt_template	str | None	Optional raw-text template with an {instruction} placeholder wrapped around the language instruction, which is lowercased and stripped before insertion (OpenVLA convention). Requires raw_text=True.	None
padding_strategy	str	HuggingFace padding strategy. "max_length" pads all sequences to max_token_len. "longest" pads to the longest sequence in the batch.	value
trust_remote_code	bool	Whether to allow tokenizers that ship custom HuggingFace code (e.g. nvidia/llama-nemotron-embed).	False

Source code in src/versatil/data/tokenization/observation_tokenizer.py

def __init__(
    self,
    tokenizer_model: str,
    observation_keys: list[str],
    bin_continuous_data: bool = True,
    num_bins: int = 256,
    max_token_len: int = 256,
    device: torch.device | None = None,
    raw_text: bool = False,
    prompt_template: str | None = None,
    padding_strategy: str = TokenPaddingStrategy.MAX_LENGTH.value,
    trust_remote_code: bool = False,
):
    """Initialize observation tokenizer.

    Args:
        tokenizer_model: HuggingFace model name (e.g., "google/gemma-2b")
        observation_keys: List of observation keys to include in prompt (order preserved)
        bin_continuous_data: Whether to bin continuous data into quantiles
        num_bins: Number of bins for quantile-based discretization
        max_token_len: Maximum token sequence length
        device: Target device for tensors
        raw_text: If True, pass language text through with only a trailing
            newline appended (no ``Task:`` prefix, no lowercasing). Use for
            VLM policies (SmolVLA, Pi0) that expect unformatted text.
        prompt_template: Optional raw-text template with an
            ``{instruction}`` placeholder wrapped around the language
            instruction, which is lowercased and stripped before insertion
            (OpenVLA convention). Requires ``raw_text=True``.
        padding_strategy: HuggingFace padding strategy. ``"max_length"``
            pads all sequences to ``max_token_len``. ``"longest"`` pads to
            the longest sequence in the batch.
        trust_remote_code: Whether to allow tokenizers that ship custom
            HuggingFace code (e.g. nvidia/llama-nemotron-embed).
    """
    if prompt_template is not None:
        if not raw_text:
            raise ValueError("prompt_template requires raw_text=True.")
        if PROMPT_TEMPLATE_INSTRUCTION_PLACEHOLDER not in prompt_template:
            raise ValueError(
                "prompt_template must contain an "
                f"'{PROMPT_TEMPLATE_INSTRUCTION_PLACEHOLDER}' placeholder, "
                f"got: {prompt_template!r}"
            )
    self.tokenizer_model = tokenizer_model
    self.observation_keys = observation_keys
    self.bin_continuous_data = bin_continuous_data
    self.num_bins = num_bins
    self.max_token_len = max_token_len
    self.device = device if device is not None else torch.device("cpu")
    self.raw_text = raw_text
    self.prompt_template = prompt_template
    self.padding_strategy = padding_strategy
    self.trust_remote_code = trust_remote_code
    self.language_tokenizer = load_huggingface_tokenizer(
        tokenizer_model=tokenizer_model, trust_remote_code=trust_remote_code
    )
    if self.language_tokenizer.pad_token is None:
        self.language_tokenizer.pad_token = self.language_tokenizer.eos_token

    self.vocab_size = len(self.language_tokenizer)
    self.binned_value_discretizers: dict[str, BinnedValueDiscretizer] = {}
    self._is_fitted = False

fit

fit(observation_data)

Fit binning tokenizers on observation data.

Parameters:

Name	Type	Description	Default
observation_data	dict[str, ndarray]	Dict mapping obs keys to arrays of shape (N, ..., D) where N is number of samples. Language keys are skipped.	required

Source code in src/versatil/data/tokenization/observation_tokenizer.py

def fit(self, observation_data: dict[str, np.ndarray]) -> None:
    """Fit binning tokenizers on observation data.

    Args:
        observation_data: Dict mapping obs keys to arrays of shape (N, ..., D)
            where N is number of samples. Language keys are skipped.
    """
    if not self.bin_continuous_data:
        self._is_fitted = True
        logging.info("Binning disabled, observation tokenizer marked as fitted")
        return

    for key in self.observation_keys:
        if key == ObsKey.LANGUAGE.value:
            continue  # Language doesn't need binning

        if key not in observation_data:
            logging.warning(
                f"Key '{key}' not found in observation data, skipping binning"
            )
            continue

        data = observation_data[key]
        discretizer = BinnedValueDiscretizer(
            num_bins=self.num_bins, device=self.device
        )
        discretizer.fit(data)
        self.binned_value_discretizers[key] = discretizer

    self._is_fitted = True
    logging.info(
        f"Fitted observation tokenizer on {len(self.observation_keys)} keys "
        f"({len(self.binned_value_discretizers)} with binning) "
        f"(model={self.tokenizer_model}, vocab_size={self.vocab_size})"
    )

tokenize

tokenize(observations, batched=False)

Tokenize observations into unified prompt.

Parameters:

Name	Type	Description	Default
observations	dict[str, Any]	Dict with observation data. - Language keys: list[str], or list[list[str]] when batched - Continuous keys: torch.Tensor or np.ndarray, with a leading batch dimension when batched	required
batched	bool	Whether inputs carry a leading batch dimension in front of the observation horizon, producing per-timestep prompts reshaped back to (B, T, seq).	False

Returns:

Type	Description
dict[str, Tensor]	Dict with: - "tokens": Token IDs (B, max_token_len) - "is_pad_observation": Padding mask (B, max_token_len)

Source code in src/versatil/data/tokenization/observation_tokenizer.py

def tokenize(
    self,
    observations: dict[str, Any],
    batched: bool = False,
) -> dict[str, torch.Tensor]:
    """Tokenize observations into unified prompt.

    Args:
        observations: Dict with observation data.
            - Language keys: list[str], or list[list[str]] when batched
            - Continuous keys: torch.Tensor or np.ndarray, with a leading
              batch dimension when batched
        batched: Whether inputs carry a leading batch dimension in front
            of the observation horizon, producing per-timestep prompts
            reshaped back to (B, T, seq).

    Returns:
        Dict with:
            - "tokens": Token IDs (B, max_token_len)
            - "is_pad_observation": Padding mask (B, max_token_len)
    """
    if not self._is_fitted:
        raise RuntimeError("Tokenizer must be fitted before encoding")

    batch_size, time_steps = None, None
    if batched:
        first_tensor = next(
            (v for v in observations.values() if isinstance(v, torch.Tensor)),
            None,
        )
        if first_tensor is not None:
            batch_size = first_tensor.shape[0]
            time_steps = first_tensor.shape[1]
        else:
            first_nested_list = next(
                (v for v in observations.values() if _is_nested_language_list(v)),
                None,
            )
            if first_nested_list is None:
                raise ValueError(
                    "Batched tokenization requires at least one tensor or "
                    "nested language list to derive the batch layout."
                )
            batch_size = len(first_nested_list)
            time_steps = len(first_nested_list[0])
        observations = {
            key: _flatten_time_dim(value=value)
            for key, value in observations.items()
        }

    prompts = self._build_prompts(observations)
    tokenized = self.language_tokenizer(
        prompts,
        padding=self.padding_strategy,
        truncation=True,
        max_length=self.max_token_len,
        return_tensors="pt",
    )
    tokens = tokenized["input_ids"]
    if "attention_mask" in tokenized:
        is_pad = ~tokenized["attention_mask"].to(torch.bool)
    else:
        is_pad = tokens == self.language_tokenizer.pad_token_id
    if batched:
        # Reshape (B*T, seq) -> (B, T, seq)
        tokens = tokens.reshape(batch_size, time_steps, -1)
        is_pad = is_pad.reshape(batch_size, time_steps, -1)

    return {
        SampleKey.TOKENIZED_OBSERVATIONS.value: tokens.to(self.device),
        SampleKey.IS_PAD_OBSERVATION.value: is_pad.to(self.device),
    }

to

to(device)

Move tokenizer to specified device.

Parameters:

Name	Type	Description	Default
device	device	Target device	required

Returns:

Type	Description
ObservationTokenizer	Self for chaining

Source code in src/versatil/data/tokenization/observation_tokenizer.py

def to(self, device: torch.device) -> "ObservationTokenizer":
    """Move tokenizer to specified device.

    Args:
        device: Target device

    Returns:
        Self for chaining
    """
    self.device = device
    for discretizer in self.binned_value_discretizers.values():
        discretizer.to(device)
    return self

state_dict

state_dict()

Get state dictionary for serialization.

Returns:

Type	Description
dict[str, Any]	Dictionary containing tokenizer state

Source code in src/versatil/data/tokenization/observation_tokenizer.py

def state_dict(self) -> dict[str, Any]:
    """Get state dictionary for serialization.

    Returns:
        Dictionary containing tokenizer state
    """
    return {
        "tokenizer_model": self.tokenizer_model,
        "observation_keys": self.observation_keys,
        "bin_continuous_data": self.bin_continuous_data,
        "num_bins": self.num_bins,
        "max_token_len": self.max_token_len,
        "vocab_size": self.vocab_size,
        "raw_text": self.raw_text,
        "prompt_template": self.prompt_template,
        "padding_strategy": self.padding_strategy,
        "trust_remote_code": self.trust_remote_code,
        "binned_value_discretizers": {
            key: discretizer.state_dict()
            for key, discretizer in self.binned_value_discretizers.items()
        },
        "is_fitted": self._is_fitted,
    }

load_state_dict

load_state_dict(state_dict)

Load state dictionary.

Parameters:

Name	Type	Description	Default
state_dict	dict[str, Any]	State dictionary from state_dict()	required

Source code in src/versatil/data/tokenization/observation_tokenizer.py

def load_state_dict(self, state_dict: dict[str, Any]) -> None:
    """Load state dictionary.

    Args:
        state_dict: State dictionary from state_dict()
    """
    self.tokenizer_model = state_dict["tokenizer_model"]
    self.observation_keys = state_dict["observation_keys"]
    self.bin_continuous_data = state_dict["bin_continuous_data"]
    self.num_bins = state_dict["num_bins"]
    self.max_token_len = state_dict["max_token_len"]
    self.vocab_size = state_dict["vocab_size"]
    self.raw_text = state_dict.get("raw_text", False)
    self.prompt_template = state_dict.get("prompt_template")
    self.padding_strategy = state_dict.get(
        "padding_strategy", TokenPaddingStrategy.MAX_LENGTH.value
    )
    self.trust_remote_code = state_dict.get("trust_remote_code", False)
    self._is_fitted = state_dict["is_fitted"]
    self.binned_value_discretizers = {}

    discretizer_states = state_dict["binned_value_discretizers"]
    for key, discretizer_state in discretizer_states.items():
        discretizer = BinnedValueDiscretizer(
            num_bins=self.num_bins, device=self.device
        )
        discretizer.load_state_dict(discretizer_state)
        self.binned_value_discretizers[key] = discretizer

save_pretrained

save_pretrained(path)

Save tokenizer to disk.

Parameters:

Name	Type	Description	Default
path	str | Path	Directory path to save tokenizer	required

Source code in src/versatil/data/tokenization/observation_tokenizer.py

def save_pretrained(self, path: str | Path) -> None:
    """Save tokenizer to disk.

    Args:
        path: Directory path to save tokenizer
    """
    path = Path(path)
    path.mkdir(parents=True, exist_ok=True)
    torch.save(self.state_dict(), path / "observation_tokenizer_state.pt")
    self.language_tokenizer.save_pretrained(path / "language_tokenizer")
    logging.info(f"Saved observation tokenizer to {path}")

from_pretrained classmethod

from_pretrained(path, device=None)

Load tokenizer from disk.

Parameters:

Name	Type	Description	Default
path	str | Path	Directory path where tokenizer was saved	required
device	device | None	Target device for tensors	None

Returns:

Type	Description
ObservationTokenizer	Loaded ObservationTokenizer instance

Source code in src/versatil/data/tokenization/observation_tokenizer.py

@classmethod
def from_pretrained(
    cls, path: str | Path, device: torch.device | None = None
) -> "ObservationTokenizer":
    """Load tokenizer from disk.

    Args:
        path: Directory path where tokenizer was saved
        device: Target device for tensors

    Returns:
        Loaded ObservationTokenizer instance
    """
    path = Path(path)
    if not path.exists():
        raise FileNotFoundError(f"Tokenizer path not found: {path}")
    state_dict = torch.load(
        path / "observation_tokenizer_state.pt",
        map_location=device or torch.device("cpu"),
        weights_only=False,
    )
    tokenizer = cls(
        tokenizer_model=state_dict["tokenizer_model"],
        observation_keys=state_dict["observation_keys"],
        bin_continuous_data=state_dict["bin_continuous_data"],
        num_bins=state_dict["num_bins"],
        max_token_len=state_dict["max_token_len"],
        device=device,
        raw_text=state_dict.get("raw_text", False),
        prompt_template=state_dict.get("prompt_template"),
        padding_strategy=state_dict.get(
            "padding_strategy", TokenPaddingStrategy.MAX_LENGTH.value
        ),
        trust_remote_code=state_dict.get("trust_remote_code", False),
    )
    tokenizer.load_state_dict(state_dict)
    tokenizer.language_tokenizer = load_huggingface_tokenizer(
        tokenizer_model=path / "language_tokenizer",
        trust_remote_code=tokenizer.trust_remote_code,
    )
    logging.info(f"Loaded observation tokenizer from {path}")
    return tokenizer