action_discretizer

action_discretizer

Discretizers for continuous action chunks.

ActionDiscretizer

Bases: ABC

Converts continuous action chunks to local discrete action IDs.

Action chunks use shape (time_horizon, action_dim). Fitting data uses shape (num_chunks, time_horizon, action_dim).

token_count abstractmethod property

token_count

Number of local discrete action IDs.

is_fitted abstractmethod property

is_fitted

Whether the discretizer can encode and decode actions.

fit abstractmethod

fit(action_chunks)

Fit on chunks with shape (num_chunks, time_horizon, action_dim).

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

@abstractmethod
def fit(self, action_chunks: np.ndarray) -> None:
    """Fit on chunks with shape (num_chunks, time_horizon, action_dim)."""

encode abstractmethod

encode(action_chunk)

Encode one chunk with shape (time_horizon, action_dim).

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

@abstractmethod
def encode(self, action_chunk: np.ndarray) -> list[int]:
    """Encode one chunk with shape (time_horizon, action_dim)."""

decode abstractmethod

decode(token_sequences)

Decode token sequences into shape (batch_size, time_horizon, action_dim).

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

@abstractmethod
def decode(self, token_sequences: list[list[int]]) -> np.ndarray:
    """Decode token sequences into shape (batch_size, time_horizon, action_dim)."""

to abstractmethod

to(device)

Move internal tensors to a device.

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

@abstractmethod
def to(self, device: torch.device) -> None:
    """Move internal tensors to a device."""

state_dict abstractmethod

state_dict()

Return serializable state.

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

@abstractmethod
def state_dict(self) -> dict[str, Any]:
    """Return serializable state."""

load_state_dict abstractmethod

load_state_dict(state_dict)

Load serializable state.

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

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

save_pretrained

save_pretrained(path)

Save optional external assets.

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

def save_pretrained(self, path: Path) -> None:
    """Save optional external assets."""
    del path

load_pretrained_assets

load_pretrained_assets(path)

Load optional external assets.

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

def load_pretrained_assets(self, path: Path) -> None:
    """Load optional external assets."""
    del path

FastActionDiscretizer

FastActionDiscretizer(use_pretrained=True, tokenizer_model='physical-intelligence/fast', time_horizon=None, action_dim=None)

Bases: ActionDiscretizer

FAST discretizer for compressed action-token sequences.

Initialize FAST processor metadata and optional pretrained assets.

Parameters:

Name	Type	Description	Default
use_pretrained	bool	Whether to use the pretrained FAST processor instead of fitting a local one.	True
tokenizer_model	str	HuggingFace model ID or local directory path.	'physical-intelligence/fast'
time_horizon	int | None	Action-chunk time horizon used to shape decoded actions. Required for pretrained processors, which skip fit; fit overwrites it for local processors.	None
action_dim	int | None	Action dimension used to shape decoded actions, with the same semantics as time_horizon.	None

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

def __init__(
    self,
    use_pretrained: bool = True,
    tokenizer_model: str = "physical-intelligence/fast",
    time_horizon: int | None = None,
    action_dim: int | None = None,
):
    """Initialize FAST processor metadata and optional pretrained assets.

    Args:
        use_pretrained: Whether to use the pretrained FAST processor
            instead of fitting a local one.
        tokenizer_model: HuggingFace model ID or local directory path.
        time_horizon: Action-chunk time horizon used to shape decoded
            actions. Required for pretrained processors, which skip fit;
            fit overwrites it for local processors.
        action_dim: Action dimension used to shape decoded actions, with
            the same semantics as ``time_horizon``.
    """
    self.use_pretrained = use_pretrained
    self.tokenizer_model = tokenizer_model
    self.processor: ProcessorMixin | None = load_fast_processor(tokenizer_model)
    self._token_count = 2048 if use_pretrained else 1024
    self._is_fitted = use_pretrained
    self.time_horizon = time_horizon
    self.action_dim = action_dim

token_count property

token_count

Number of local discrete action IDs.

is_fitted property

is_fitted

Whether the FAST processor can encode and decode actions.

fit

fit(action_chunks)

Fit a local FAST processor on shape (num_chunks, time_horizon, action_dim).

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

def fit(self, action_chunks: np.ndarray) -> None:
    """Fit a local FAST processor on shape (num_chunks, time_horizon, action_dim)."""
    if self.use_pretrained:
        raise ValueError(
            "Cannot fit a pretrained FAST action discretizer. "
            "Set use_pretrained=False to fit FAST on local data."
        )
    if self.processor is None:
        raise RuntimeError("FAST processor not initialized")
    self.time_horizon = action_chunks.shape[1]
    self.action_dim = action_chunks.shape[2]
    self.processor = self.processor.fit(
        action_chunks,
        time_horizon=self.time_horizon,
        action_dim=self.action_dim,
    )
    self._is_fitted = True

encode

encode(action_chunk)

Encode one chunk with shape (time_horizon, action_dim).

The chunk shape only seeds the decode shape when it is still unknown: padded chunks arrive with their padded rows already dropped, so a known time horizon must never be overwritten by a shorter chunk.

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

def encode(self, action_chunk: np.ndarray) -> list[int]:
    """Encode one chunk with shape (time_horizon, action_dim).

    The chunk shape only seeds the decode shape when it is still unknown:
    padded chunks arrive with their padded rows already dropped, so a
    known time horizon must never be overwritten by a shorter chunk.
    """
    if self.processor is None:
        raise RuntimeError("FAST processor not initialized")
    if self.time_horizon is None:
        self.time_horizon = action_chunk.shape[-2]
    if self.action_dim is None:
        self.action_dim = action_chunk.shape[-1]
    return self.processor(action_chunk)[0]

decode

decode(token_sequences)

Decode FAST tokens into shape (batch_size, time_horizon, action_dim).

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

def decode(self, token_sequences: list[list[int]]) -> np.ndarray:
    """Decode FAST tokens into shape (batch_size, time_horizon, action_dim)."""
    if self.processor is None:
        raise RuntimeError("FAST processor not initialized")
    if self.time_horizon is None or self.action_dim is None:
        raise RuntimeError(
            "FAST action discretizer shape is unknown; encode or load a fitted "
            "discretizer before decoding"
        )
    decoded_actions = [
        self._decode_fast_token_sequence(
            token_sequence=sequence,
        )
        for sequence in token_sequences
    ]
    return np.stack(decoded_actions)

to

to(device)

No-op device transfer for the processor-backed discretizer.

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

def to(self, device: torch.device) -> None:
    """No-op device transfer for the processor-backed discretizer."""
    del device

state_dict

state_dict()

Return serializable FAST discretizer state.

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

def state_dict(self) -> dict[str, Any]:
    """Return serializable FAST discretizer state."""
    return {
        "type": ActionDiscretizerType.FAST.value,
        "use_pretrained": self.use_pretrained,
        "tokenizer_model": self.tokenizer_model,
        "token_count": self.token_count,
        "is_fitted": self.is_fitted,
        "time_horizon": self.time_horizon,
        "action_dim": self.action_dim,
    }

load_state_dict

load_state_dict(state_dict)

Load FAST discretizer state.

States saved before the chunk shape was persisted carry None for time_horizon/action_dim; those keep any values already set on this instance instead of erasing them.

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

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

    States saved before the chunk shape was persisted carry None for
    ``time_horizon``/``action_dim``; those keep any values already set on
    this instance instead of erasing them.
    """
    self.use_pretrained = state_dict["use_pretrained"]
    self.tokenizer_model = state_dict.get("tokenizer_model", self.tokenizer_model)
    self._token_count = state_dict.get("token_count", self._token_count)
    self._is_fitted = state_dict["is_fitted"]
    loaded_time_horizon = state_dict.get("time_horizon")
    if loaded_time_horizon is not None:
        self.time_horizon = loaded_time_horizon
    loaded_action_dim = state_dict.get("action_dim")
    if loaded_action_dim is not None:
        self.action_dim = loaded_action_dim

save_pretrained

save_pretrained(path)

Save fitted local FAST processor assets.

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

def save_pretrained(self, path: Path) -> None:
    """Save fitted local FAST processor assets."""
    if self.processor is not None and not self.use_pretrained:
        self.processor.save_pretrained(str(path / "fast_processor"))

load_pretrained_assets

load_pretrained_assets(path)

Load saved local FAST processor assets when present.

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

def load_pretrained_assets(self, path: Path) -> None:
    """Load saved local FAST processor assets when present."""
    fast_path = path / "fast_processor"
    if fast_path.exists():
        self.processor = load_fast_processor(str(fast_path))
        self._is_fitted = True

BinnedActionDiscretizer

BinnedActionDiscretizer(num_bins=256, device=None, binning_strategy=value, min_value=-1.0, max_value=1.0)

Bases: ActionDiscretizer

Per-value quantile binning for chunks with shape (time_horizon, action_dim).

Initialize per-value binning for action chunks.

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

def __init__(
    self,
    num_bins: int = 256,
    device: torch.device | None = None,
    binning_strategy: str = BinningStrategy.UNIFORM.value,
    min_value: float = -1.0,
    max_value: float = 1.0,
):
    """Initialize per-value binning for action chunks."""
    self.binner = BinnedValueDiscretizer(
        num_bins=num_bins,
        device=device,
        binning_strategy=binning_strategy,
        min_value=min_value,
        max_value=max_value,
    )
    self.time_horizon: int | None = None
    self.action_dim: int | None = None

token_count property

token_count

Number of bins used for each action value.

is_fitted property

is_fitted

Whether bin edges have been fitted.

fit

fit(action_chunks)

Fit bin edges from shape (num_chunks, time_horizon, action_dim).

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

def fit(self, action_chunks: np.ndarray) -> None:
    """Fit bin edges from shape (num_chunks, time_horizon, action_dim)."""
    self.time_horizon = action_chunks.shape[1]
    self.action_dim = action_chunks.shape[2]
    self.binner.fit(action_chunks)

encode

encode(action_chunk)

Encode one chunk with shape (time_horizon, action_dim).

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

def encode(self, action_chunk: np.ndarray) -> list[int]:
    """Encode one chunk with shape (time_horizon, action_dim)."""
    tokens = self.binner.encode(action_chunk)
    return tokens.reshape(-1).detach().cpu().tolist()

decode

decode(token_sequences)

Decode bin-ID sequences into shape (batch_size, time_horizon, action_dim).

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

def decode(self, token_sequences: list[list[int]]) -> np.ndarray:
    """Decode bin-ID sequences into shape (batch_size, time_horizon, action_dim)."""
    if self.time_horizon is None or self.action_dim is None:
        raise RuntimeError("Binned action discretizer shape is unknown")

    expected_len = self.time_horizon * self.action_dim
    normalized_sequences = []
    for sequence in token_sequences:
        if len(sequence) != expected_len:
            raise ValueError(
                "Binned action token sequence has invalid length: "
                f"expected {expected_len} tokens for shape "
                f"({self.time_horizon}, {self.action_dim}), got {len(sequence)}."
            )
        token_array = np.asarray(sequence)
        if np.any(token_array < 0) or np.any(token_array >= self.token_count):
            raise ValueError(
                "Binned action token sequence contains IDs outside the valid "
                f"range [0, {self.token_count - 1}]."
            )
        normalized_sequences.append(token_array.astype(np.int64).tolist())

    tokens = torch.tensor(normalized_sequences, dtype=torch.long)
    tokens = tokens.reshape(
        len(token_sequences), self.time_horizon, self.action_dim
    )
    return self.binner.decode(tokens).detach().cpu().numpy()

to

to(device)

Move bin-edge tensors to a device.

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

def to(self, device: torch.device) -> None:
    """Move bin-edge tensors to a device."""
    self.binner.to(device)

state_dict

state_dict()

Return serializable binned discretizer state.

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

def state_dict(self) -> dict[str, Any]:
    """Return serializable binned discretizer state."""
    return {
        "type": ActionDiscretizerType.BINNED.value,
        "num_bins": self.token_count,
        "time_horizon": self.time_horizon,
        "action_dim": self.action_dim,
        "binner": self.binner.state_dict(),
        "is_fitted": self.is_fitted,
    }

load_state_dict

load_state_dict(state_dict)

Load binned discretizer state.

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

def load_state_dict(self, state_dict: dict[str, Any]) -> None:
    """Load binned discretizer state."""
    self.time_horizon = state_dict.get("time_horizon")
    self.action_dim = state_dict.get("action_dim")
    self.binner.load_state_dict(state_dict["binner"])