backup: sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/config.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cbade0cbb2e1b288670272377ec8a11e727628777ca6cbbb86a5eacac668e4e9
+size 8170

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/full_config.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:99892c7c03fe07d35cbc77302b4d97ec4ef4b8f065666da110a0cada336dd3f8
+size 16031

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/git.branch

@@ -0,0 +1 @@
+petko/bridge_all_annotations

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/git.hash

@@ -0,0 +1 @@
+e1c6fda68772ca50f04ab2935360fadb1ea63e2c

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/hf_export/keen-fuchsia-mandrill/src/common_pizero_fm_qwen3_vl.py

@@ -0,0 +1,570 @@
+from functools import cached_property
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type
+
+import torch
+import torch.nn.attention.flex_attention
+import transformers
+import transformers.models.qwen3_vl.modeling_qwen3_vl
+from backports.strenum import StrEnum
+from databib.dataclasses import Dataclass, dataclass
+from databib.dataclasses.dataclass import DataclassT
+from databib.utils.classproperty import classproperty
+
+
+class ReferenceFrame(StrEnum):
+    """
+    Indicates the frame frame w.r.t. which translation or rotation is expressed.
+    Note that each of translation and rotation has its own (possibly different) ReferenceFrame value.
+
+    WORLD: Only for completeness, not yet used. Will become relevant when navigation is introduced.
+    ROBOT_BASE: Translation/rotation expressed in absolute robot base frame
+    ROBOT_BASE_DELTA:
+        - Translation expressed as delta value w.r.t. the previous EEF translation pose
+            The delta value is defined in the robot base frame (rather than in the current EEF frame)
+        - Rotation expressed as w.r.t. the previous rotation pose
+            The axis of rotation is defined in the robot base frame (rather than in the current EEF frame)
+    ROBOT_BASE_RELATIVE: Same as ROBOT_BASE_DELTA, but the sequence is expressed w.r.t.the 0-th element
+        instead of the previous element
+    EEF: Translation/rotation expressed in the current end-effector frame
+    EEF_DELTA:
+        - Translation expressed as delta value w.r.t. the previous EEF translation pose
+            The delta value is defined in the current EEF frame (rather than in the robot base frame)
+        - Rotation expressed as w.r.t. the previous rotation pose
+            The axis of rotation is defined in the current EEF frame (rather than in the robot base frame)
+    """
+
+    ROBOT_BASE = 'robot_base'
+    ROBOT_BASE_DELTA = 'robot_base_delta'
+    ROBOT_BASE_RELATIVE = 'robot_base_relative'
+    EEF_RELATIVE = EEF = 'eef_relative'
+    EEF_DELTA = 'eef_delta'
+    CAMERA = 'camera'
+    UNKNOWN = 'unknown'
+
+    @classproperty
+    def robot_frames(cls) -> set['ReferenceFrame']:
+        return {
+            ReferenceFrame.ROBOT_BASE,
+            ReferenceFrame.ROBOT_BASE_DELTA,
+            ReferenceFrame.ROBOT_BASE_RELATIVE,
+        }
+
+    @classproperty
+    def eef_frames(cls) -> set['ReferenceFrame']:
+        return {ReferenceFrame.EEF, ReferenceFrame.EEF_RELATIVE, ReferenceFrame.EEF_DELTA}
+
+    @classproperty
+    def delta_frames(cls) -> set['ReferenceFrame']:
+        return {ReferenceFrame.ROBOT_BASE_DELTA, ReferenceFrame.EEF_DELTA}
+
+    @classproperty
+    def relative_frames(cls) -> set['ReferenceFrame']:
+        return {ReferenceFrame.ROBOT_BASE_RELATIVE, ReferenceFrame.EEF_RELATIVE}
+
+    @classproperty
+    def core_frames(cls) -> set['ReferenceFrame']:
+        return {ReferenceFrame.ROBOT_BASE, ReferenceFrame.EEF}
+
+    def to_relative(self) -> 'ReferenceFrame':
+        if self in self.robot_frames:
+            return self.ROBOT_BASE_RELATIVE
+        if self in self.eef_frames:
+            return self.EEF_RELATIVE
+        raise ValueError(f'Cannot convert frame {self} to relative frame')
+
+    def to_delta(self) -> 'ReferenceFrame':
+        if self in self.robot_frames:
+            return self.ROBOT_BASE_DELTA
+        if self in self.eef_frames:
+            return self.EEF_DELTA
+        raise ValueError(f'Cannot convert frame {self} to delta frame')
+
+    def to_core(self) -> 'ReferenceFrame':
+        if self in self.robot_frames:
+            return self.ROBOT_BASE
+        if self in self.eef_frames:
+            return self.EEF
+        raise ValueError(f'Cannot convert frame {self} to relative frame')
+
+
+class RotationFormat(StrEnum):
+    """Determines how rotations will be encoded in the loaded batch"""
+
+    EULER = 'euler'
+    QUATERNION = 'quaternion'
+    ROTMAT = 'rotmat'
+
+
+class ResizeMode(StrEnum):
+    """
+    Different modes for resizing images.
+    """
+
+    MATCH_WIDTH = 'match_width'
+    MATCH_HEIGHT = 'match_height'
+    MATCH_MAX = 'match_max'
+    NAIVE = 'naive'
+    SMART = 'smart'
+    PAD = 'pad'
+    CROP = 'crop'
+
+
+def expand_dims(tensor: torch.Tensor, ndim: int, order: Sequence[int]) -> torch.Tensor:
+    """
+    Expand the dimensions of `tensor` to `ndim` such that all new dimensions have size of 1
+    Args:
+        tensor: torch.Tensor of any shape
+        ndim: Number of output dimensions. Must be >= `tensor.ndim`
+        order: Sequence of size `tensor.ndim + 1`. Contains only values of 1 and a single value of -1,
+            indicating where the new `ndim - tensor.ndim` dimensions will be inserted
+    Returns:
+        torch.Tensor with dimensions `ndim`, a view of `tensor`
+
+    Ex:
+        expand_dims(torch.ones([2, 3, 4]), ndim=5, order=[1, -1, 1, 1]).shape -> [2, 1, 1, 3, 4]
+        expand_dims(torch.ones([2, 3, 4]), ndim=5, order=[-1, 1, 1, 1]).shape -> [1, 1, 2, 3, 4]
+        expand_dims(torch.ones([2, 3, 4]), ndim=5, order=[1, 1, 1, -1]).shape -> [2, 3, 4, 1, 1]
+    """
+    assert tensor.ndim <= ndim, f'{tensor.ndim} > {ndim}; shape={tensor.shape}'
+    assert len(order) == tensor.ndim + 1, f'{len(order)} != {tensor.ndim + 1}; shape={tensor.shape}'
+    order = list(order)
+    assert order.count(-1) == 1, 'Order must have exactly one value of -1'
+    assert order.count(1) == len(order) - 1, 'Order must have exactly len(order) - 1 values of 1'
+    if tensor.ndim == ndim:
+        return tensor
+    insert_index = order.index(-1)
+    view = list(tensor.shape[:insert_index]) + [1] * (ndim - tensor.ndim) + list(tensor.shape[insert_index:])
+    tensor = tensor.view(view)
+    return tensor
+
+
+def compare_dicts(dict_0: Dict[str, Any], dict_1: Dict[str, Any], comparison_function: Callable) -> bool:
+    if set(dict_0.keys()) != set(dict_1.keys()):
+        return False
+    for key, _ in dict_0.items():
+        if type(dict_0[key]) != type(dict_1[key]):
+            return False
+        if isinstance(dict_0[key], dict):
+            result = compare_dicts(dict_0[key], dict_1[key], comparison_function)
+        else:
+            result = comparison_function(dict_0[key], dict_1[key])
+            if isinstance(result, torch.Tensor):
+                result = bool(result.all())
+        if not result:
+            return False
+    return True
+
+
+def tensor_size_bytes(tensor: Optional[torch.Tensor]) -> int:
+    if tensor is None:
+        return 0
+    if not isinstance(tensor, torch.Tensor):
+        raise RuntimeError('Provided data is not a torch.Tensor: ', tensor)
+    bytes_per_element = tensor.element_size()
+    return bytes_per_element * tensor.numel()
+
+
+def tensor_dataclass(cls: Type[DataclassT], **kwargs) -> Type[DataclassT]:
+    cls = dataclass(cls, eq=False, **kwargs)
+    return cls
+
+
+@tensor_dataclass
+class TensorDataclass(Dataclass):
+    """
+    Extends Dataclass with common torch.Tensor utilities.
+        - Can contain non-tensor fields, but some member functions might ignore these fields
+            or explicitly raise errors.
+        - Useful for packing batches, input and output data for ML models
+        - When using for input / output data for ML models, it's recommended to keep only torch.Tensor
+            fields to allow for supporting functionality such as torch.jit.script
+    """
+
+    def __eq__(self, other) -> bool:
+        if type(other) is not type(self):
+            return False
+        return compare_dicts(self.as_json(), other.as_json(), lambda x, y: x == y)
+
+    def __ne__(self, other) -> bool:
+        return not self == other
+
+    def __hash__(self):
+        raise ValueError(f'Hash function not implemented for {self.__class__.__name__}.')
+
+    def calc_size_bytes(self) -> int:
+        return sum(
+            (
+                tensor_size_bytes(value)
+                for (_, value) in self.items(recursive=True)
+                if isinstance(value, torch.Tensor)
+            )
+        )
+
+    def calc_size_megabytes(self) -> float:
+        return self.calc_size_bytes() / 2**20
+
+    def cpu(self) -> 'TensorDataclass':
+        return self.to(device='cpu')
+
+    def to(self, *, device=None, dtype=None, copy=False, non_blocking=False) -> 'TensorDataclass':
+        assert device is not None or dtype is not None
+        return self.apply(
+            lambda value: value.to(device=device, dtype=dtype, copy=copy, non_blocking=non_blocking)
+            if isinstance(value, torch.Tensor)
+            else value
+        )
+
+    def float32(self) -> 'TensorDataclass':
+        return self.apply(
+            lambda value: value.to(dtype=torch.float32)
+            if isinstance(value, torch.Tensor) and value.dtype.is_floating_point
+            else value
+        )
+
+    def detach(self) -> 'TensorDataclass':
+        return self.apply(lambda value: value.detach() if isinstance(value, torch.Tensor) else value)
+
+    def __getitem__(self, index) -> 'TensorDataclass':
+        def extract(obj):
+            if obj is None:
+                return None
+            if isinstance(obj, torch.Tensor):
+                return obj[index]
+            raise ValueError(f'Cannot slice {obj.__class__.__name__} object')
+
+        return self.apply(extract)
+
+    @property
+    def device(self) -> Optional[torch.device]:
+        """
+        Returns the device on which tensors in this dataclass reside. If tensors are on
+        different devices, raises RuntimeError. If no tensors in the class, returns None
+        """
+        devices = [
+            value.device
+            for (key, value) in self.items()
+            if isinstance(value, (TensorDataclass, torch.Tensor))
+        ]
+        devices = [d for d in devices if d is not None]
+        if len(devices) == 0:
+            return None
+        if len(set(devices)) == 1:
+            return devices[0]
+        (key, device) = (None, None)
+        for k, value in self.items():
+            if value is None:
+                continue
+            if device is None:
+                device = value.device
+                key = k
+            elif device != value.device:
+                raise RuntimeError(
+                    f'Inconsistent device for instance of {self.__class__.__name__}. Device of field {key} is {device}, while device of field {k} is {value.device}'
+                )
+        raise RuntimeError
+
+    def to_shared_memory(self) -> 'TensorDataclass':
+        """Move all tensors in the dataclass to shared memory"""
+        return self.apply(lambda value: value.share_memory_() if isinstance(value, torch.Tensor) else value)
+
+    def pin_memory(self) -> 'TensorDataclass':
+        """Used for pinning memory during dataloading. Do not modify the name of the function"""
+        return self.apply(lambda value: value.pin_memory() if isinstance(value, torch.Tensor) else value)
+
+
+@tensor_dataclass
+class ModelTarget(TensorDataclass):
+    """
+    Only relevant for supervised learning.
+    Packs regression / classification target values that we input in the loss
+    """
+
+
+@tensor_dataclass
+class RoboticsTarget(ModelTarget):
+    control_tokens_ids: Optional[torch.Tensor]
+    text_tokens_ids: Optional[torch.Tensor]
+    translation: torch.Tensor
+    rotation: torch.Tensor
+    gripper: torch.Tensor
+    valid_mask: torch.Tensor
+
+
+@tensor_dataclass
+class PolicyControlPlan(TensorDataclass):
+    """
+    Abstraction class relevant for control tasks. Note that `ModelOutput` might not contain the actual
+    controls we want to use on the robot in the environment. Examples:
+        - `ModelOutput` contains logits, since computing losses on logits is more numerically stable.
+            We need to convert these logits to actual controls for the actual robot
+        - `ModelOutput` contains an entire costmap from which we need to extract waypoints
+        - `ModelOutput` contains unnormalized quaternion or rotation matrix that need to be normalized
+        - `ModelOutput` contains 2D/3D positions from which we need to extract speed and steering
+    `PolicyControlPlan`
+        - Extracts actual physical representation from `ModelOutput` that we can use to dervie the controls
+        - Doesn't necessarily contain the controls themselves, but they can be derived from this data
+        - **Interpretable control plan which we can visualize, interpret and compare to the real data**
+            - Ex: Controls might be in speed and steering, but we likely want to compare 2D/3D positions
+                instead of controls for metrics and visualizations
+            - Ex: Robot control is usually a single timestep, while `PolicyControlPlan` contains
+                controls over multiple timesteps
+        - Can have different abstractions, e.g.
+            - End effector 3D translation and rotation (positional control)
+            - Speed and steering for a vehicle (actuator control)
+            - 3D waypoints for a path to be followed
+        - Usually **unnormalized** values into physical units (vs normalized `ModelOutput`)
+    Main purpose: (Human) Interpretable control plans and metadata that can be used for visualization,
+        metrics and debugging
+    """
+
+
+@tensor_dataclass
+class RoboticsControlPlan(PolicyControlPlan):
+    translation_m: torch.Tensor
+    rotmat: torch.Tensor
+    gripper_prob: torch.Tensor
+    valid_mask: torch.Tensor
+
+    def __post_init__(self):
+        super().__post_init__()
+        assert self.translation_m.ndim == 3, self.translation_m.shape
+        assert self.rotmat.ndim == 3, self.rotmat.shape
+        assert self.gripper_prob.ndim == 3, self.gripper_prob.shape
+
+
+@tensor_dataclass
+class ModelOutput(TensorDataclass):
+    """
+    Packs data which an NN model outputs. Note this can contain a lot of metadata
+    such as intermediate outputs, probabilities, visualizations, etc
+    In the case of robot control, the action class is not guaranteed to be part of this
+    class, but we must be able to derive an action from the data in this class
+    """
+
+
+@tensor_dataclass
+class RoboticsInput(TensorDataclass):
+    images: Dict[str, torch.Tensor]
+    input_ids: torch.Tensor
+    attn_mask: torch.Tensor
+    ee_pose_translation: torch.Tensor
+    ee_pose_rotation: torch.Tensor
+    gripper: torch.Tensor
+    joints: torch.Tensor
+    control_tokens_ids: Optional[torch.Tensor]
+
+    @property
+    def inputs_embeds(self) -> Optional[torch.Tensor]:
+        return None
+
+    @property
+    def past_key_values(self) -> Optional[List[torch.Tensor]]:
+        return None
+
+    @cached_property
+    def multimodal_indices(self) -> torch.Tensor:
+        """
+        Returns a torch.Tensor containing only the indices of the batch examples which are multimodal.
+        Return shape is [B]
+        """
+        return torch.arange(self.input_ids.shape[0], dtype=torch.int64, device=self.input_ids.device)
+
+    @cached_property
+    def unimodal_indices(self) -> torch.Tensor:
+        """
+        Returns a torch.Tensor containing only the indices of the batch examples which are unimodal.
+        Return shape is [B]
+        """
+        return torch.tensor([], dtype=torch.int64, device=self.input_ids.device)
+
+
+@tensor_dataclass
+class FlowInput(TensorDataclass):
+    timestep: torch.Tensor
+    translation_t: torch.Tensor
+    rotation_t: torch.Tensor
+    gripper_t: torch.Tensor
+    translation_t0: torch.Tensor
+    rotation_t0: torch.Tensor
+    gripper_t0: torch.Tensor
+
+
+@tensor_dataclass
+class RoboticsFlowInput(RoboticsInput):
+    """Input to the entire Robotics VLM"""
+
+    flow_input: FlowInput
+
+
+@tensor_dataclass
+class DiffusionInput(TensorDataclass):
+    timestep: torch.Tensor
+    noised_translation: torch.Tensor
+    noised_rotation: torch.Tensor
+    noised_gripper: torch.Tensor
+
+
+@tensor_dataclass
+class LLMOutput(TensorDataclass):
+    """Fork of transformers.modeling_outputs.CausalLMOutputWithPast"""
+
+    input_ids: torch.Tensor
+    logits: Optional[torch.Tensor]
+    output_ids: Optional[torch.Tensor]
+    loss: Optional[torch.Tensor]
+    past_key_values: List[Tuple[torch.Tensor, torch.Tensor]]
+    hidden_states: List[torch.Tensor]
+    text_mask: torch.Tensor
+    image_mask: torch.Tensor
+
+    @classmethod
+    def from_transformers(
+        cls,
+        input_ids: torch.Tensor,
+        llm_output: transformers.modeling_outputs.CausalLMOutputWithPast,
+        text_mask: torch.Tensor,
+        image_mask: torch.Tensor,
+    ) -> 'LLMOutput':
+        return LLMOutput(
+            input_ids=input_ids,
+            logits=getattr(llm_output, 'logits', None),
+            output_ids=None,
+            loss=getattr(llm_output, 'loss', None),
+            past_key_values=list(llm_output.past_key_values)
+            if llm_output.past_key_values is not None
+            else [],
+            hidden_states=list(llm_output.hidden_states) if llm_output.hidden_states is not None else [],
+            text_mask=text_mask,
+            image_mask=image_mask,
+        )
+
+    def compress(self, ignore_index: int = -100) -> 'LLMOutput':
+        """
+        Compress the data contained in the class so it can be moved between CPU and GPU or concatenated
+        much faster:
+            - hidden_states - huge tensors; take a lot of CPU time to move across devices or concat
+            - past_key_values - huge tensors; take a lot of CPU time to move across devices or concat
+            - logits - huge last dimension; takes a lot of CPU time to move across devices or concat
+        """
+        replace: Dict[str, Any] = {'hidden_states': [], 'past_key_values': [], 'loss': None}
+        if self.logits is not None:
+            replace['logits'] = None
+            if self.output_ids is None:
+                assert (
+                    self.text_mask is not None
+                ), 'text_mask is required to compute output_ids when output_ids is None'
+                assert (
+                    self.logits.shape[:2] == self.text_mask.shape
+                ), 'logits and text_mask batch and sequence dimensions must match to compute output_ids'
+                predicted_ids = self.logits.argmax(dim=-1)
+                output_ids = torch.where(self.text_mask, predicted_ids, ignore_index)
+                replace['output_ids'] = output_ids
+        return self.replace(**replace)
+
+
+@tensor_dataclass
+class RoboticsOutput(ModelOutput):
+    translation: Optional[torch.Tensor]
+    rotation: Optional[torch.Tensor]
+    gripper: Optional[torch.Tensor]
+    token_logits: Optional[torch.Tensor]
+    token_ids: Optional[torch.Tensor]
+    llm_output: LLMOutput
+
+    def compress(self, ignore_index: int = -100) -> 'RoboticsOutput':
+        """
+        Compress output and drop unnecessary components to speed up transfer GPU <-> CPU.
+        Note that LLM logits can be extremely expensive since their size is [B, S, vocab_size], which
+        can reach millions or billions of values for large vocab_size
+        """
+        replace: Dict[str, Any] = {
+            'llm_output': self.llm_output.compress(ignore_index=ignore_index),
+            'token_logits': None,
+        }
+        if self.token_logits is not None and self.token_ids is None:
+            replace['token_ids'] = torch.argmax(self.token_logits, dim=-1)
+        return self.replace(**replace)
+
+
+@tensor_dataclass
+class VLMOutput(TensorDataclass):
+    llm_output: LLMOutput
+    vit_tokens: Optional[torch.Tensor]
+    attn_mask: torch.Tensor
+
+    def compress(self, ignore_index: int = -100) -> 'VLMOutput':
+        """
+        Compress output and drop unnecessary components to speed up transfer GPU <-> CPU.
+        Note that LLM logits can be extremely expensive since their size is [B, S, vocab_size], which
+        can reach millions or billions of values for large vocab_size
+        """
+        return self.replace(llm_output=self.llm_output.compress(ignore_index=ignore_index))
+
+
+def is_quaternion(quaternion: torch.Tensor) -> bool:
+    return quaternion.shape[-1] == 4
+
+
+def quaternion_half_cover(quaternion: torch.Tensor) -> torch.Tensor:
+    """
+    Flip quaternions so they cover only a half the space. If the q_w is negative, flip the quaternion.
+    If q_w is 0, then choose such that the first non-zero component is positive. Note that geometrically,
+    this doesn't correspond to a single hemisphere of the unit sphere. Follows
+    https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.transform.Rotation.as_quat.html#scipy.spatial.transform.Rotation.as_quat
+    """
+    assert is_quaternion(quaternion), quaternion.shape
+    with torch.no_grad():
+        is_zero = quaternion == 0
+        flip_condition = (
+            (quaternion[..., -1:] < 0)
+            | is_zero[..., -1:] & (quaternion[..., 0:1] < 0)
+            | is_zero[..., -1:] & is_zero[..., 0:1] & (quaternion[..., 1:2] < 0)
+            | is_zero[..., -1:] & is_zero[..., 0:1] & is_zero[..., 1:2] & (quaternion[..., 2:3] < 0)
+        )
+    quaternion = torch.where(flip_condition, -quaternion, quaternion)
+    return quaternion
+
+
+def is_rotmat_3x3(rotmat: torch.Tensor) -> bool:
+    return rotmat.shape[-2:] == torch.Size([3, 3])
+
+
+def is_rotmat_9(rotmat: torch.Tensor) -> bool:
+    return rotmat.shape[-1] == 9
+
+
+def rotmat_as_9(rotmat: torch.Tensor) -> torch.Tensor:
+    """Convert any rotmat input to [..., 9] shape"""
+    if is_rotmat_9(rotmat):
+        return rotmat
+    if is_rotmat_3x3(rotmat):
+        return rotmat.reshape(*rotmat.shape[:-2], 9)
+    raise ValueError(f"Can't convert tensor of shape {rotmat.shape} to a 3x3 rotation matrix")
+
+
+def is_rotmat(rotmat: torch.Tensor) -> bool:
+    """
+    Checks if the tensor shape matches that of a rotmat. However, it's not guaranteed the data is a
+    valid rotmat. `is_orthonormal_rotmat` performs this additional check.
+    NOTE: This might incorrectly return True if the underlying data is euler angles and accidentally
+    `rotmat.shape[-2:] == [3, 3]`. This would happen very rarely, but use with caution
+    """
+    return is_rotmat_3x3(rotmat) or is_rotmat_9(rotmat)
+
+
+def rotmat_as_3x3(rotmat: torch.Tensor) -> torch.Tensor:
+    """Convert any rotmat input to [..., 3, 3] shape"""
+    if rotmat.shape[-1] == 9:
+        return rotmat.reshape(*rotmat.shape[:-1], 3, 3)
+    if rotmat.shape[-2:] == torch.Size([3, 3]):
+        return rotmat
+    raise ValueError(f"Can't convert tensor of shape {rotmat.shape} to a 3x3 rotation matrix")
+
+
+def rotmat_inverse(rotation: torch.Tensor) -> torch.Tensor:
+    assert is_rotmat(rotation), f'Expected a rotation matrix, but got shape {rotation.shape}'
+    rotmat = rotmat_as_3x3(rotation)
+    rotmat = rotmat.transpose(-1, -2)
+    if is_rotmat_9(rotation):
+        rotmat = rotmat_as_9(rotmat)
+    return rotmat

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/hf_export/keen-fuchsia-mandrill/src/configuration_pizero_fm_qwen3_vl.py

@@ -0,0 +1,330 @@
+from typing import Any, Dict, List, Optional
+
+from databib.config import Config
+
+from .common_pizero_fm_qwen3_vl import ReferenceFrame, ResizeMode, RotationFormat
+
+
+class ConfigurableModuleConfig(Config):
+    @property
+    def pretrained(self) -> bool:
+        return not self.pretrain_config.empty
+
+
+class FourierFeaturesProjectorConfig(ConfigurableModuleConfig):
+    in_features: int
+    num_features: int = 256
+    layers: List[int] = [256, 512, 256]
+    activation: str = 'GELU'
+    norm: Optional[str] = None
+
+
+class RotaryPositionalEncodingConfig(ConfigurableModuleConfig):
+    num_embeddings: int
+    embedding_dim: int
+    base: int = 10000
+    cached: bool = True
+
+
+class PiZeroFlowMatchingDecoderBlockConfig(ConfigurableModuleConfig):
+    feature_size: int
+    head_dim: int = 128
+    num_heads: int = 32
+    num_kv_heads: int = 1
+    hidden_size: int
+    activation: str = 'GELU'
+    activation_kwargs: Dict[str, Any] = {}
+    norm: str = 'RMSNorm'
+    dropout: float = 0.0
+    attn_implementation: str = 'sdpa'
+    position_embed_config: RotaryPositionalEncodingConfig
+
+
+class PiZeroFlowMatchingDecoderConfig(ConfigurableModuleConfig):
+    num_blocks: int
+    block_config: PiZeroFlowMatchingDecoderBlockConfig
+
+
+class RobotStateProjectorConfig(ConfigurableModuleConfig):
+    layers: List[int] = []
+    mode: str
+    activation: str = 'GELU'
+    fourier: bool = False
+
+    def __post_init__(self):
+        super().__post_init__()
+        assert self.mode in [
+            'ee_pose',
+            'ee_pose_gripper',
+            'ee_pose_joints',
+            'joints',
+            'all',
+            'none',
+        ], self.mode
+
+
+class FourierFeaturesConfig(ConfigurableModuleConfig):
+    num_features: int = 256
+    learnable_features: bool = False
+    max_period: float = 10000.0
+    layers: List[int] = [256, 512, 256]
+    activation: str = 'SiLU'
+    norm: Optional[str] = None
+
+
+class NoisedControlProjectorConfig(ConfigurableModuleConfig):
+    time_embed: FourierFeaturesConfig
+    layers: List[int] = []
+    activation: str = 'SiLU'
+    norm: Optional[str] = None
+
+
+class PiZeroFlowMatchingModuleConfig(ConfigurableModuleConfig):
+    token_size: int = 1024
+    noised_control_proj_config: NoisedControlProjectorConfig
+    robot_state_proj_config: RobotStateProjectorConfig
+    control_decoder_config: PiZeroFlowMatchingDecoderConfig
+    rotation_components: int = 3
+
+
+class VLMConfig(ConfigurableModuleConfig):
+    pass
+
+
+class InputSequencingConfig(Config):
+    """
+    past_frames_sequence_length: number of past images needed in a single robot state
+    past_scalars_sequence_length: number of past scalar state data, e.g. actions, poses, etc,
+        needed in a single robot state
+    past_frames_stride_sec: sampling rate, determines how far apart in time each point in the sequence
+        is. If None, ignored and takes the default data collection frequency from the dataset
+    past_scalars_stride_sec: similar to past_frames_stride_sec
+
+    sequence_frames: number of temporally-sequential points in a single example in the batch
+    sequence_frames_stride_sec: sampling rate
+
+    Understanding sequence_frames:
+        TODO: sequences are possibly useful in some rare cases, maybe sequence modeling problems,
+            but yet to be confirmed. Keeping for now, but could be removed if proved unnecessary
+
+        - past_scalars_sequence_length, past_frames_sequence_length, future_controls_sequence_length,
+            future_frames_sequence_length are hyperparameters refering to a SINGLE dataset example / 'state'.
+            It is assumed that `past_scalars_sequence_length` and `past_frames_sequence_length` are the min
+            number of observations that comprise a single 'state'
+        - sequence_frames is a hyperparameter refering to the entire learning process. It controls the size
+            of the sequence dimension in the batch. It's treated similarly to the batch dimension, with the
+            difference that points in the sequence dimensions are temporally aligned. Unlike `past_*`
+            attributes, in supervised learning a label is loaded for every point in the sequence dimension
+            and the loss usually computed over the entire sequence dimension.
+    """
+
+    past_scalars_sequence_length: int = 1
+    past_frames_sequence_length: int = 1
+    past_scalars_stride_sec: Optional[float] = None
+    past_frames_stride_sec: Optional[float] = None
+    sequence_frames: int = 1
+    sequence_frames_stride_sec: Optional[float] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        assert self.past_scalars_sequence_length >= 1, self.past_scalars_sequence_length
+        assert self.past_frames_sequence_length >= 1, self.past_frames_sequence_length
+        assert self.sequence_frames >= 1, self.sequence_frames
+        if self.past_frames_stride_sec is not None:
+            assert self.past_frames_stride_sec >= 0.0, self.past_frames_stride_sec
+        if self.past_scalars_stride_sec is not None:
+            assert self.past_scalars_stride_sec >= 0.0, self.past_scalars_stride_sec
+        if self.sequence_frames_stride_sec is not None:
+            assert self.sequence_frames_stride_sec >= 0.0, self.sequence_frames_stride_sec
+
+    def assert_same_past(self) -> None:
+        assert (
+            self.past_frames_stride_sec == self.past_scalars_stride_sec
+        ), f'{self.past_frames_stride_sec} != {self.past_scalars_stride_sec}'
+        assert (
+            self.past_frames_sequence_length == self.past_scalars_sequence_length
+        ), f'{self.past_frames_sequence_length} != {self.past_scalars_sequence_length}'
+
+
+class OutputSequencingConfig(Config):
+    """
+    future_controls_sequence_length: number of control steps in the future the model predicts
+    future_frames_sequence_length: number of future frames the model predicts
+        (only relevant for neural networks that learn some sort of a world model)
+
+    future_controls_sequence_stride_sec / future_frames_sequence_stride_sec: sampling rate
+        that determines how far apart in time each point in the sequence is. If None,
+        ignored and takes the default data collection frequency from the dataset
+
+    future_control_offset_sec: time interval between the last observation and the first
+    point at which control is predicted. Serves as a 'causality hyperparameter', allowing
+    for predicting controls slightly further into the future in environments with dynamics
+    where the observed effects of an action appear slightly later
+    """
+
+    future_controls_sequence_length: int = 1
+    future_controls_sequence_stride_sec: Optional[float] = None
+    future_frames_sequence_length: int = 1
+    future_frames_sequence_stride_sec: Optional[float] = None
+    future_control_offset_sec: float = 0.0
+
+    def __post_init__(self):
+        super().__post_init__()
+        assert self.future_controls_sequence_length >= 1, self.future_controls_sequence_length
+        assert self.future_frames_sequence_length >= 1, self.future_frames_sequence_length
+        assert self.future_control_offset_sec >= 0.0, self.future_control_offset_sec
+        if self.future_controls_sequence_stride_sec is not None:
+            assert self.future_controls_sequence_stride_sec >= 0.0, self.future_controls_sequence_stride_sec
+        if self.future_frames_sequence_stride_sec is not None:
+            assert self.future_frames_sequence_stride_sec >= 0.0, self.future_frames_sequence_stride_sec
+
+
+class ControlDataIOConfig(InputSequencingConfig, OutputSequencingConfig):
+    pass
+
+
+class NormalizerConfig(Config):
+    pass
+
+
+class RotationStereomapNormalizerConfig(NormalizerConfig):
+    factor: float
+
+
+class IdentityNormalizerConfig(NormalizerConfig):
+    pass
+
+
+class DatasetStatsNormalizerConfig(NormalizerConfig):
+    stats_filepath: str
+    stats_key: str = ''
+    component_name: str
+    mode: str
+
+    def __post_init__(self):
+        super().__post_init__()
+        assert self.mode in {'mean', 'bounds', 'bounds_q99'}, self.mode
+
+
+class BoundsNormalizerConfig(NormalizerConfig):
+    low: List[float]
+    high: List[float]
+
+    def __post_init__(self):
+        super().__post_init__()
+        if len(self.low) != len(self.high):
+            raise ValueError(
+                f'Low and high bounds must have the same length, but got {self.low} and {self.high}'
+            )
+        for low, high in zip(self.low, self.high, strict=True):
+            assert low < high, f'Low bound {low} must be less than high bound {high}'
+
+
+class ControlTokenizerConfig(Config):
+    pass
+
+
+class EmptyTokenizerConfig(ControlTokenizerConfig):
+    pass
+
+
+class VLAMProcessorConfig(Config):
+    control_io_config: ControlDataIOConfig
+    joints_obs_norm: BoundsNormalizerConfig
+    translation_obs_norm: DatasetStatsNormalizerConfig
+    rotation_obs_norm: IdentityNormalizerConfig
+    translation_control_norm: BoundsNormalizerConfig
+    rotation_control_norm: RotationStereomapNormalizerConfig
+    translation_obs_frame: ReferenceFrame = ReferenceFrame.ROBOT_BASE
+    rotation_obs_frame: ReferenceFrame = ReferenceFrame.ROBOT_BASE
+    translation_control_frame: ReferenceFrame = ReferenceFrame.ROBOT_BASE_DELTA
+    rotation_control_frame: ReferenceFrame = ReferenceFrame.EEF_DELTA
+    rotation_format: RotationFormat
+    image_resize: ResizeMode = ResizeMode.SMART
+    control_tokenizer_config: EmptyTokenizerConfig
+
+    def __post_init__(self):
+        super().__post_init__()
+        if (
+            self.rotation_obs_frame != ReferenceFrame.ROBOT_BASE
+            or self.translation_obs_frame != ReferenceFrame.ROBOT_BASE
+        ):
+            raise NotImplementedError()
+
+    @property
+    def delta_controls(self) -> bool:
+        translation_is_delta = self.translation_control_frame in (
+            ReferenceFrame.ROBOT_BASE_DELTA,
+            ReferenceFrame.EEF_DELTA,
+        )
+        rotation_is_delta = self.rotation_control_frame in (
+            ReferenceFrame.ROBOT_BASE_DELTA,
+            ReferenceFrame.EEF_DELTA,
+        )
+        if translation_is_delta != rotation_is_delta:
+            raise NotImplementedError(
+                'Delta controls for only one of translation or rotation not yet supported'
+            )
+        return translation_is_delta
+
+
+class RegressionProcessorConfig(VLAMProcessorConfig):
+    pass
+
+
+class PiZeroFlowProcessorConfig(RegressionProcessorConfig):
+    num_inference_steps: int
+    r0_distribution: str = 'uniform'
+    timestep_distribution: str
+    distribution_hyperparams: Dict[str, Any] = {}
+    sig_min: float = 0.001
+
+    def __post_init__(self):
+        super().__post_init__()
+        assert self.r0_distribution in ['normal', 'uniform']
+        if self.rotation_obs_frame != ReferenceFrame.ROBOT_BASE:
+            raise NotImplementedError()
+
+
+class VLMProcessorConfig(Config):
+    pass
+
+
+class ImageSizeConfig(Config):
+    width: int
+    height: int
+
+
+class Qwen3VLProcessorConfig(VLMProcessorConfig):
+    image_sizes: Dict[str, ImageSizeConfig] = {'main': ImageSizeConfig(width=256, height=256)}
+
+
+class Qwen3VLConfig(VLMConfig):
+    """
+    VLM config for Qwen3-VL model.
+
+    Attributes:
+    model_id: The identifier of the pre-trained Qwen3-VL model to be used
+    attn_implementation: The attention implementation to be used in the model
+    processor_config: Configuration for the VLM processor
+    lm_head: If True, includes the language model head in the model; otherwise, it replaces
+    it with an identity layer. It helps to save memory when the LM head is not needed.
+    mixed_modality_forward: If True, replaces the default forward method of Qwen3-VL
+    model with a custom one that can handle mixed modality inputs, including text-only
+    inputs.
+    """
+
+    model_id: str = 'Qwen/Qwen3-VL-2B-Instruct'
+    attn_implementation: str = 'flash_attention_2'
+    processor_config: Qwen3VLProcessorConfig
+    lm_head: bool = True
+    mixed_modality_forward: bool = True
+
+
+class VLAMConfig(ConfigurableModuleConfig):
+    processor_config: PiZeroFlowProcessorConfig
+    vlm_config: Qwen3VLConfig
+    control_module_config: PiZeroFlowMatchingModuleConfig
+
+
+MainModelConfig = VLAMConfig

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/hf_export/keen-fuchsia-mandrill/src/format.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:743eb01140fcf46b321de8b04dd211859f85464c3ed40e94ef47b095f266c3f1
+size 5766

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/hf_export/keen-fuchsia-mandrill/src/model_config.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a76bce2911e91a10be67ae8e69620c8c175c45498a97592b1e7c273f0f5ae906
+size 2673

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/hf_export/keen-fuchsia-mandrill/src/modeling_pizero_fm_qwen3_vl.py

@@ -0,0 +1,2067 @@
+import collections
+import math
+import warnings
+from abc import abstractmethod
+from functools import cached_property
+from typing import Any, Callable, Dict, List, Optional, Protocol, Tuple, Type, TypeVar, Union
+
+import PIL.Image
+import roma
+import torch
+import torch.nn.attention.flex_attention
+import transformers
+import transformers.models.qwen3_vl.modeling_qwen3_vl
+from databib.config import Configurable
+from databib.template import Template
+
+from .common_pizero_fm_qwen3_vl import (
+    DiffusionInput,
+    FlowInput,
+    LLMOutput,
+    RoboticsFlowInput,
+    RoboticsInput,
+    RoboticsOutput,
+    RotationFormat,
+    VLMOutput,
+    expand_dims,
+    is_quaternion,
+    is_rotmat,
+    is_rotmat_3x3,
+    quaternion_half_cover,
+    rotmat_as_3x3,
+    rotmat_as_9,
+    rotmat_inverse,
+)
+from .configuration_pizero_fm_qwen3_vl import (
+    ConfigurableModuleConfig,
+    FourierFeaturesConfig,
+    FourierFeaturesProjectorConfig,
+    ImageSizeConfig,
+    NoisedControlProjectorConfig,
+    PiZeroFlowMatchingDecoderBlockConfig,
+    PiZeroFlowMatchingDecoderConfig,
+    PiZeroFlowMatchingModuleConfig,
+    Qwen3VLConfig,
+    Qwen3VLProcessorConfig,
+    RobotStateProjectorConfig,
+    RotaryPositionalEncodingConfig,
+    VLAMConfig,
+    VLMConfig,
+)
+from .processing_pizero_fm_qwen3_vl import EmptyTokenizer, PiZeroFlowMatchingProcessor, VLMProcessor
+
+
+class GemmaRMSNorm(torch.nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-06):
+        super().__init__()
+        self.eps = eps
+        self.weight = torch.nn.Parameter(torch.zeros(dim))
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float())
+        output = output * (1.0 + self.weight.float())
+        return output.type_as(x)
+
+    def extra_repr(self):
+        return f'{tuple(self.weight.shape)}, eps={self.eps}'
+
+
+class Qwen3VLProcessor(VLMProcessor[Qwen3VLProcessorConfig]):
+    def __init__(self, config: Qwen3VLProcessorConfig, hf_processor: transformers.AutoProcessor):
+        super().__init__(config)
+        self.hf_processor = hf_processor
+        self.turn_start_token = '<|im_start|>'
+        self.turn_end_token = '<|im_end|>'
+        self.assistant_header = 'assistant'
+        self.sep_token = '\n'
+        self.turn_start_id = self.hf_processor.tokenizer.added_tokens_encoder[self.turn_start_token]
+        self.turn_end_id = self.hf_processor.tokenizer.added_tokens_encoder[self.turn_end_token]
+        self.assistant_header_id = self.hf_processor.tokenizer('assistant')['input_ids'][0]
+        self.sep_id = self.hf_processor.tokenizer('\n')['input_ids'][0]
+
+    @property
+    def tokenizer(self) -> transformers.PreTrainedTokenizerBase:
+        return self.hf_processor.tokenizer
+
+    @property
+    def image_sizes(self) -> Dict[str, ImageSizeConfig]:
+        return self.config.image_sizes
+
+    @cached_property
+    def _flattened_patch_dim(self) -> int:
+        """
+        Return the dimensionality (number of scalar elements) of a flattened image patch.
+        This is computed as (patch_size ** 2) * 3 * merge_size, where:
+        - patch_size is the side length (in pixels) of a square patch from hf_processor.image_processor.patch_size,
+        - 3 corresponds to RGB channels,
+        - merge_size is hf_processor.image_processor.merge_size - for more info refer to the Qwen3-VL paper and code.
+
+        Example:
+            For patch_size=16 and merge_size=2 this returns 16 * 16 * 3 * 2 == 1536.
+
+        Returns:
+            int: Number of values in a flattened patch (used as the per-patch input dimension for visual embeddings).
+        """
+        return (
+            self.hf_processor.image_processor.patch_size**2
+            * 3
+            * self.hf_processor.image_processor.merge_size
+        )
+
+    @cached_property
+    def num_image_patches(self) -> Dict[int, int]:
+        hf_image_processor = self.hf_processor.image_processor
+        num_image_patches_per_camera = {}
+        for camera_name, camera_image_size in self.image_sizes.items():
+            (width, height) = (camera_image_size.width, camera_image_size.height)
+            num_image_patches_per_camera[camera_name] = hf_image_processor.get_number_of_image_patches(
+                width, height, {}
+            )
+        return num_image_patches_per_camera
+
+    @cached_property
+    def num_image_tokens(self) -> Dict[str, int]:
+        """
+        Number of image tokens per camera
+
+        Returns:
+            Dict[str, int]: number of image tokens per camera
+        """
+        return {
+            camera_name: num_image_patches // self.hf_processor.image_processor.merge_size**2
+            for (camera_name, num_image_patches) in self.num_image_patches.items()
+        }
+
+    def preprocess_inputs(
+        self, chat: List[str], images: Dict[str, List[PIL.Image.Image]]
+    ) -> Dict[str, torch.Tensor | Dict[str, torch.Tensor]]:
+        for key, value in images.items():
+            if not isinstance(value, list):
+                raise TypeError(f'Camera {key} contains values of type {type(value)} instead of list')
+            if len(value) > 1:
+                raise NotImplementedError(
+                    f'Multiple images per camera not supported yet, but camera {key} contains {len(value)} images'
+                )
+        messages = []
+        for i, text in enumerate(chat):
+            if i % 2 == 0:
+                content = []
+                if i == 0:
+                    for _, camera_images in images.items():
+                        content.append({'type': 'image', 'image': camera_images[0]})
+                content.append({'type': 'text', 'text': text})
+                messages.append({'role': 'user', 'content': content})
+            else:
+                content = [{'type': 'text', 'text': text}]
+                messages.append({'role': 'assistant', 'content': content})
+        hf_inputs = self.hf_processor.apply_chat_template(
+            messages, add_generation_prompt=False, tokenize=True, return_dict=True, return_tensors='pt'
+        )
+        turn_end_idxs = torch.nonzero(hf_inputs.input_ids[0] == self.turn_end_id).squeeze(1).tolist()
+        target_ids = hf_inputs.input_ids.clone()
+        next_turn_end_idx = 0
+        start_message_idx = 0
+        for msg in messages:
+            if next_turn_end_idx < len(turn_end_idxs):
+                end_message_idx = turn_end_idxs[next_turn_end_idx] + 1
+            else:
+                end_message_idx = hf_inputs.input_ids.shape[1] - 1
+            if msg['role'] == 'user':
+                target_ids[0, start_message_idx : end_message_idx + 1] = self.ignore_index
+            elif msg['role'] == 'assistant':
+                target_ids[0, start_message_idx : start_message_idx + 3] = self.ignore_index
+                target_ids[0, end_message_idx - 1 : end_message_idx + 1] = self.ignore_index
+            else:
+                raise ValueError('Unknown role')
+            start_message_idx = end_message_idx + 1
+            next_turn_end_idx += 1
+        input_ids = hf_inputs.input_ids.squeeze(0)
+        target_ids = target_ids.squeeze(0)
+        attn_mask = hf_inputs.attention_mask.squeeze(0)
+        images = {'pixel_values': hf_inputs.pixel_values, 'image_grid_thw': hf_inputs.image_grid_thw}
+        return {'input_ids': input_ids, 'target_ids': target_ids, 'images': images, 'attn_mask': attn_mask}
+
+    @property
+    def ignore_index(self) -> int:
+        return -100
+
+
+ConfigurableModuleConfigT = TypeVar('ConfigurableModuleConfigT', bound=ConfigurableModuleConfig)
+
+
+class ConfigurableModule(
+    torch.nn.Module, Configurable[ConfigurableModuleConfigT], Template[ConfigurableModuleConfigT]
+):
+    """
+    Helper base class that inherits from both torch.nn.Module and Configurable.
+    Provides `PretrainedModuleConfig()` functionality safely and out of the box
+    """
+
+    def __init__(self, config: ConfigurableModuleConfigT):
+        Configurable[self.ConfigT].__init__(self, config)
+        torch.nn.Module.__init__(self)
+
+
+def make_mlp(
+    layer_sizes: List[int],
+    activation: str | Type[torch.nn.Module],
+    norm: str | Type[torch.nn.Module] | None = torch.nn.LayerNorm,
+    activate_final: bool = False,
+    bias: bool = True,
+) -> torch.nn.Sequential:
+    """
+    Args:
+        layer_sizes: List of layer sizes. The first value is the number of input features and the last
+            value is the number of output features
+        activation: str or the class of the activation. If str, it should be the exact name of
+            the activation module under torch.nn, e.g. 'ReLU', 'SiLU', 'GeLU'. Use 'Identity' if
+            no activation wanted
+        norm: type of normalization. Same type as `activation`. Ex: `torch.nn.LayerNorm`, 'LayerNorm', etc
+    """
+    if len(layer_sizes) == 0:
+        return torch.nn.Identity()
+    assert len(layer_sizes) > 1, 'Need to provide input and output layer sizes at least'
+    if isinstance(activation, str):
+        TorchActivation: Type[torch.nn.Module] = getattr(torch.nn, activation)
+    else:
+        TorchActivation: Type[torch.nn.Module] = activation
+    assert issubclass(TorchActivation, torch.nn.Module), TorchActivation
+    if isinstance(norm, str):
+        TorchNorm: Type[torch.nn.Module] = getattr(torch.nn, norm)
+    elif norm is None:
+        TorchNorm: Type[torch.nn.Module] = torch.nn.Identity
+    else:
+        TorchNorm: Type[torch.nn.Module] = norm
+    assert issubclass(TorchNorm, torch.nn.Module), TorchNorm
+
+    def make_norm_act(modules: dict[str, torch.nn.Module], empty: bool):
+        return {} if empty else modules
+
+    module = torch.nn.Sequential(
+        *[
+            torch.nn.Sequential(
+                collections.OrderedDict(
+                    {
+                        'linear': torch.nn.Linear(in_features, out_features, bias=bias),
+                        **make_norm_act(
+                            {'norm': TorchNorm(out_features), 'act': TorchActivation()},
+                            empty=i == len(layer_sizes) - 2 and not activate_final,
+                        ),
+                    }
+                )
+            )
+            for (i, (in_features, out_features)) in enumerate(
+                zip(layer_sizes[:-1], layer_sizes[1:], strict=True)
+            )
+        ]
+    )
+    return module
+
+
+class FourierFeaturesProjector(ConfigurableModule[FourierFeaturesProjectorConfig]):
+    def __init__(self, config: FourierFeaturesProjectorConfig):
+        super().__init__(config)
+        self.feature_proj = torch.nn.Linear(
+            in_features=self.config.in_features, out_features=self.config.num_features // 2, bias=False
+        )
+        self.layers: torch.nn.Sequential = make_mlp(
+            self.config.layers, activation=self.config.activation, norm=self.config.norm, activate_final=False
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Compute Fourier features and project them via MLP
+        Args:
+            x: Input tensor of shape [..., in_features]
+        Returns:
+            torch.Tensor: Fourier features of shape [..., out_features]
+        """
+        frequencies = 2 * math.pi * self.feature_proj(x)
+        output = torch.cat([torch.cos(frequencies), torch.sin(frequencies)], dim=-1)
+        output = self.layers(output)
+        return output
+
+
+class RobotStateProjector(ConfigurableModule[RobotStateProjectorConfig]):
+    """Pack robot state and project to a single token per timestep"""
+
+    def __init__(self, config: RobotStateProjectorConfig):
+        super().__init__(config)
+        if self.config.fourier:
+            self.robot_state_tokens_proj = FourierFeaturesProjector(
+                FourierFeaturesProjectorConfig(
+                    in_features=self.config.layers[0],
+                    num_features=self.config.layers[1],
+                    layers=self.config.layers[1:],
+                    activation=self.config.activation,
+                )
+            )
+        else:
+            self.robot_state_tokens_proj = make_mlp(
+                layer_sizes=self.config.layers, activation=self.config.activation, norm=torch.nn.LayerNorm
+            )
+
+    def forward(self, inputs: RoboticsInput) -> Optional[torch.Tensor]:
+        """
+        Returns:
+            torch.Tensor of shape [B, num_past_steps, token_size] or None (if mode == 'none')
+        """
+        if self.config.mode == 'ee_pose':
+            robot_state = torch.cat([inputs.ee_pose_translation, inputs.ee_pose_rotation], dim=-1)
+        elif self.config.mode == 'ee_pose_gripper':
+            robot_state = torch.cat(
+                [inputs.ee_pose_translation, inputs.ee_pose_rotation, inputs.gripper], dim=-1
+            )
+        elif self.config.mode == 'ee_pose_joints':
+            robot_state = torch.cat(
+                [inputs.ee_pose_translation, inputs.ee_pose_rotation, inputs.joints], dim=-1
+            )
+        elif self.config.mode == 'joints':
+            robot_state = inputs.joints
+        elif self.config.mode == 'all':
+            robot_state = torch.cat(
+                [inputs.ee_pose_translation, inputs.ee_pose_rotation, inputs.gripper, inputs.joints], dim=-1
+            )
+        elif self.config.mode == 'none':
+            robot_state = torch.tensor([], device=inputs.ee_pose_translation.device).view(
+                inputs.ee_pose_translation.shape[0],
+                0,
+                self.config.layers[0] if len(self.config.layers) > 0 else 0,
+            )
+        else:
+            raise NotImplementedError(f'Unknown image tokens mode {self.config.mode}')
+        output = self.robot_state_tokens_proj(robot_state)
+        return output
+
+
+class FourierFeatures(ConfigurableModule[FourierFeaturesConfig]):
+    def __init__(self, config: FourierFeaturesConfig):
+        super().__init__(config)
+        if self.config.learnable_features:
+            self.linear = torch.nn.Linear(
+                in_features=1, out_features=self.config.num_features // 2, bias=False
+            )
+        else:
+            half_dim = self.config.num_features // 2
+            freqs = torch.log(torch.tensor(self.config.max_period)) / (half_dim - 1)
+            freqs = torch.exp(-freqs * torch.arange(half_dim))
+            self.register_buffer('freqs', freqs)
+        self.layers: torch.nn.Sequential = make_mlp(
+            self.config.layers, activation=self.config.activation, norm=self.config.norm, activate_final=False
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Compute Fourier features and project them via MLP
+        Args:
+            x: Input tensor of shape [..., 1]
+        Returns:
+            torch.Tensor: Fourier features of shape [..., num_features] or [..., layers[-1]]
+        """
+        assert x.shape[-1] == 1 and x.ndim > 1, x.shape
+        if self.config.learnable_features:
+            frequencies = 2 * math.pi * self.linear(x)
+        else:
+            frequencies = x * expand_dims(self.freqs, x.ndim, [-1, 1])
+        output = torch.cat([torch.cos(frequencies), torch.sin(frequencies)], dim=-1)
+        output = self.layers(output)
+        return output
+
+
+class NoisedControlProjector(ConfigurableModule[NoisedControlProjectorConfig]):
+    """Pack noised control (translation, rotation, gripper) and project to a single token per timestep"""
+
+    def __init__(self, config: NoisedControlProjectorConfig):
+        super().__init__(config)
+        self.input_projector = torch.nn.Linear(
+            in_features=self.config.layers[0], out_features=self.config.layers[1] // 2, bias=False
+        )
+        self.time_embed = FourierFeatures(self.config.time_embed)
+        self.layers = make_mlp(
+            self.config.layers[1:],
+            activation=self.config.activation,
+            norm=self.config.norm,
+            activate_final=False,
+            bias=False,
+        )
+
+    def forward(self, inputs: FlowInput | DiffusionInput) -> Optional[torch.Tensor]:
+        """
+        Returns:
+            torch.Tensor of shape [B, num_control_timesteps, token_size]
+        """
+        noised_controls = torch.cat([inputs.translation_t, inputs.rotation_t, inputs.gripper_t], dim=-1)
+        noised_controls = self.input_projector(noised_controls)
+        timestep = self.time_embed(inputs.timestep)
+        timestep = timestep.expand(-1, noised_controls.shape[1], -1)
+        features = torch.cat([timestep, noised_controls], dim=-1)
+        output = self.layers(features)
+        return output
+
+
+def make_position_indices(
+    position_indices: Optional[torch.Tensor],
+    seq_length: int,
+    device: torch.device,
+    max_seq_length: Optional[int],
+) -> torch.Tensor:
+    if position_indices is not None:
+        position_indices = position_indices.to(dtype=torch.int64)
+    else:
+        position_indices = torch.arange(seq_length, dtype=torch.int64, device=device).view(1, -1)
+    if not torch.max(position_indices) < max_seq_length:
+        raise IndexError(
+            f'position_indices={position_indices} contains index out of bounds of num_embeddings={max_seq_length}'
+        )
+    return position_indices
+
+
+class RotaryPositionalEncoding(ConfigurableModule[RotaryPositionalEncodingConfig]):
+    """
+    Rotary Positional Embeddings (RoPE) from https://arxiv.org/abs/2104.09864
+    Reference implementations:
+        - https://github.com/meta-llama/llama/blob/main/llama/model.py#L80
+        - transformers.models.gemma.modeling_gemma.GemmaRotaryEmbedding
+        - transformers.models.llama.modeling_llama.LlamaRotaryEmbedding
+
+    If cached=True, we cache the embeddings for each position up to `num_embeddings`
+    """
+
+    def __init__(self, config: RotaryPositionalEncodingConfig):
+        super().__init__(config)
+        inv_freq = 1.0 / self.config.base ** (
+            torch.arange(0, self.config.embedding_dim, 2, dtype=torch.float32) / self.config.embedding_dim
+        )
+        self.register_buffer('inv_freq', inv_freq, persistent=False)
+        self._build_cache()
+
+    def _build_cache(self) -> None:
+        if not self.config.cached:
+            return
+        position_indices = torch.arange(self.config.num_embeddings, dtype=torch.float32)
+        indices_inv_freq = torch.einsum('i, j -> ij', position_indices, self.inv_freq)
+        sin = torch.sin(indices_inv_freq)
+        cos = torch.cos(indices_inv_freq)
+        self.register_buffer('sin_cache', sin, persistent=False)
+        self.register_buffer('cos_cache', cos, persistent=False)
+
+    def forward(
+        self, tokens: torch.Tensor, position_indices: Optional[torch.Tensor] = None, apply: bool = True
+    ) -> torch.Tensor:
+        """
+        Args:
+            tokens: torch.Tensor of shape [B, ..., S, head_dim], where `...` might be any number of dims
+            position_indices: torch.Tensor of shape [B | 1, S]. The indices of tokens within the sequence
+            apply: If True, apply the positional embedding on tokens and return the result
+        Returns:
+            torch.Tensor of the same shape as `tokens` with positional embedding applied on tokens
+        """
+        assert apply, f'{self.__class__} does not support applying embeddings externally'
+        position_indices = make_position_indices(
+            position_indices,
+            seq_length=tokens.shape[-2],
+            device=tokens.device,
+            max_seq_length=self.config.num_embeddings,
+        )
+        if self.config.cached:
+            sin = self.sin_cache[position_indices]
+            cos = self.cos_cache[position_indices]
+            sin = torch.cat([sin, sin], dim=-1)
+            cos = torch.cat([cos, cos], dim=-1)
+        else:
+            inv_freq = self.inv_freq.view(1, -1, 1).to(dtype=torch.float32)
+            position_indices = position_indices.to(dtype=torch.float32).unsqueeze(1)
+            with warnings.catch_warnings():
+                warnings.filterwarnings(
+                    'ignore',
+                    message='In CPU autocast, but the target dtype is not supported. Disabling autocast.',
+                )
+                with torch.autocast(device_type=tokens.device.type, dtype=torch.float32):
+                    freqs = (inv_freq @ position_indices).transpose(1, 2)
+                    emb = torch.cat((freqs, freqs), dim=-1)
+                    (sin, cos) = (torch.sin(emb), torch.cos(emb))
+        (sin, cos) = (sin.to(dtype=tokens.dtype), cos.to(dtype=tokens.dtype))
+        sin = expand_dims(sin, tokens.ndim, order=[1, -1, 1, 1])
+        cos = expand_dims(cos, tokens.ndim, order=[1, -1, 1, 1])
+        tokens = tokens * cos + self._rotate_invert_half(tokens) * sin
+        return tokens
+
+    @staticmethod
+    def _rotate_invert_half(x: torch.Tensor) -> torch.Tensor:
+        x1 = x[..., : x.shape[-1] // 2]
+        x2 = x[..., x.shape[-1] // 2 :]
+        return torch.cat((-x2, x1), dim=-1)
+
+
+EAGER_ATTN = 'eager'
+
+FLEX_ATTN = 'flex'
+
+SDPA_ATTN = 'sdpa'
+
+FLASH_ATTN = 'flash_attention_2'
+
+
+def flash_attn_2_forward(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    attn_mask: Optional[torch.Tensor],
+    dropout: float,
+    is_causal: bool,
+    **kwargs,
+):
+    """
+    Applies flash attention 2 on already linearly projected query, key and value.
+
+    Args:
+        query_states: Linearly projected query embedding of shape [B, num_heads, L, head_dim]
+        key_states: Linearly projected key embedding of shape [B, num_kv_heads, S, head_dim]
+        value_states: Linearly projected value embedding of shape [B, num_kv_heads, S, head_dim]
+        attn_mask: dtype torch.bool and shape [B, S].
+            If bool, False values indicate masked positions (opposite of sdpa_attn)
+            If attn_mask is None, full-bidirectional attention or causal attention is used depdening
+            on the value of `is_causal`.
+            NOTE: Doesn't support 4D attn_mask, unlike sdpa_attn
+        num_heads: Number of heads for query
+        num_kv_heads: Number of heads for keys and values
+        is_training: True if running in training mode, False otherwise
+        dropout: Dropout probability applied to attention weights
+        is_causal: If True, apply additional causal masking whe computing attention
+    Returns:
+        Tuple with entries:
+            - Attention block output: torch.Tensor of shape [B, L, num_heads, head_dim]
+            - None
+    """
+    del kwargs
+    assert (
+        attn_mask is None or attn_mask.ndim == 2 and attn_mask.dtype == torch.bool
+    ), f'{FLASH_ATTN} supports only bool attn_mask of shape [B, S] or None'
+    query_states = query_states.transpose(1, 2)
+    key_states = key_states.transpose(1, 2)
+    value_states = value_states.transpose(1, 2)
+    raise NotImplementedError('Correctness not yet confirmed')
+    attn_output = transformers.modeling_flash_attention_utils._flash_attention_forward(
+        query_states=query_states,
+        key_states=key_states,
+        value_states=value_states,
+        attention_mask=attn_mask,
+        query_length=query_states.shape[1],
+        position_ids=None,
+        dropout=dropout,
+        sliding_window=None,
+        use_top_left_mask=False,
+        is_causal=is_causal,
+        deterministic=True,
+    )
+    return attn_output, None
+
+
+def is_full_attn(attn_mask: Optional[torch.Tensor]) -> bool:
+    """
+    Return True if attn_mask doesn't contain any masked out positions, False otherwise
+    """
+    if attn_mask is None:
+        return True
+    if attn_mask.dtype == torch.bool:
+        return torch.all(attn_mask == 1).item()
+    if attn_mask.dtype.is_floating_point:
+        return torch.all(attn_mask == 0).item()
+    raise TypeError(f'Unrecognized dtype {attn_mask.dtype}')
+
+
+def unmask_unattended(attn_mask: torch.Tensor, mask_value: Optional[float] = None) -> torch.Tensor:
+    """
+    Copy-pased from `transformers.modeling_attn_mask_utils.AttentionMaskConverter._unmask_unattended`
+
+    Attend to all tokens in fully-masked rows. This is required by F.scaled_dot_product_attention
+    memory-efficient attention path. Otherwise, results are NaN
+    Details: https://github.com/pytorch/pytorch/issues/110213
+
+    Args:
+        attn_mask: [B, 1 | num_heads, query_seq_len, kv_seq_len] or [B, query_seq_len, kv_seq_len], float dtype
+        mask_value: The value inside `attn_mask` that corresponds to masked elements
+    Returns:
+
+    For example, if `attn_mask` is (e.g. here left-padding case)
+    ```
+    [
+        [[
+            [0, 0, 0],
+            [0, 0, 0],
+            [0, 0, 1]
+        ]],
+        [[
+            [1, 0, 0],
+            [1, 1, 0],
+            [1, 1, 1]
+        ]],
+        [[
+            [0, 0, 0],
+            [0, 1, 0],
+            [0, 1, 1]
+        ]]
+    ]
+    ```
+    then the modified `attn_mask` will be
+    ```
+    [
+        [[
+            [1, 1, 1],   <-- modified
+            [1, 1, 1],   <-- modified
+            [0, 0, 1]
+        ]],
+        [[
+            [1, 0, 0],
+            [1, 1, 0],
+            [1, 1, 1]
+        ]],
+        [[
+            [1, 1, 1],   <-- modified
+            [0, 1, 0],
+            [0, 1, 1]
+        ]]
+    ]
+    ```
+    """
+    assert attn_mask.dtype.is_floating_point, attn_mask.dtype
+    if mask_value is None:
+        mask_value = torch.finfo(attn_mask.dtype).min
+    return attn_mask * ~torch.all(attn_mask == mask_value, dim=-1, keepdim=True)
+
+
+@torch.no_grad()
+def attn_mask_to_float(attn_mask: torch.Tensor, dtype: Optional[torch.dtype] = None) -> torch.Tensor:
+    """
+    Convert a 4D mask of type bool to `dtype`. If the attn_mask isn't 4D or isn't bool, raise error
+    """
+    assert attn_mask.ndim == 4, attn_mask.shape
+    assert attn_mask.dtype == torch.bool, attn_mask.dtype
+    if dtype is None:
+        dtype = torch.get_autocast_dtype(attn_mask.device.type)
+    mask_value = torch.finfo(dtype).min
+    attn_mask = torch.zeros(attn_mask.shape, dtype=dtype, device=attn_mask.device).masked_fill(
+        ~attn_mask, mask_value
+    )
+    attn_mask = unmask_unattended(attn_mask, mask_value)
+    return attn_mask
+
+
+@torch.no_grad()
+def make_4d_float_attn_mask(
+    attn_mask: Optional[torch.Tensor],
+    query_seq_length: int,
+    kv_seq_length: int,
+    dtype: torch.dtype,
+    device: torch.device,
+    batch_size: int,
+) -> torch.Tensor:
+    """
+    Creates a 4D mask of shape [B | 1, 1, query_length, kv_seq_length] from a 2D mask of shape [B, kv_seq_length].
+    If the input `attn_mask` is already 4D: if dtype=torch.bool, convert to dtype, else do nothing
+    If the input is None, output is a full bi-directional attn_mask
+
+    Args:
+        attn_mask: A 2D attention mask of shape [B, kv_seq_length] or [B, 1, query_length, kv_seq_length]
+            and dtype bool. False values indicate masked out positions
+        query_seq_length: The query sequence length (L)
+        kv_seq_length: The key-value sequence length (S). When `transformers.StaticCache` is used, this should
+            equal the cache size to account for zero-padding the part of the cache that is not yet filled.
+        dtype: Output dtype
+        device: Output device
+        batch_size: Batch size
+    Returns:
+        torch.Tensor of shape [B | 1, 1, query_length, kv_seq_length] (i.e. [B | 1, 1, L, S]).
+            Contains zero at unmasked positions and `torch.finfo(dtype).min` at masked positions
+    """
+    if attn_mask is not None and attn_mask.ndim == 4:
+        if attn_mask.dtype == torch.bool:
+            attn_mask = attn_mask_to_float(attn_mask, dtype=dtype)
+        elif attn_mask.dtype != dtype:
+            raise TypeError(f'Expected attn_mask.dtype={dtype}, but got {attn_mask.dtype}')
+        return attn_mask
+    mask_value = torch.finfo(dtype).min
+    output_mask = torch.zeros([batch_size, 1, query_seq_length, kv_seq_length], dtype=dtype, device=device)
+    if attn_mask is not None:
+        assert attn_mask.dtype == torch.bool, f'Unsupported dtype {attn_mask.dtype}'
+        mask_length = attn_mask.shape[-1]
+        if mask_length != kv_seq_length:
+            raise NotImplementedError(f'{mask_length} != {kv_seq_length} not properly supported yet')
+        inverted_mask = ~attn_mask.view(batch_size, 1, 1, mask_length)
+        output_mask[..., :mask_length] = output_mask[..., :mask_length].masked_fill(inverted_mask, mask_value)
+    return output_mask
+
+
+@torch.no_grad()
+def make_attn_mask_causal(attn_mask: torch.Tensor, cache_position: torch.Tensor) -> torch.Tensor:
+    """
+    Args:
+        attn_mask: 4D tensor of shape [B | 1, 1, query_seq_len, kv_seq_len] (i.e. [B | 1, 1, L, S]) of float
+            dtype (NOT bool!). Masked positions contain the value `torch.finfo(dtype).min`
+        cache_position: torch.Tensor of type torch.int64 and shape [query_seq_len]. Contained values
+            are index positions of the query tokens in the sequence. During training, this would usually
+            be torch.arange(query_seq_len), but during generate, this would usually be a tensor sequence
+            with 1 element indicating the position of the token currently being generated
+    Returns:
+        torch.Tensor of the same shape as attn_mask. Contains zero at unmasked positions and
+            `torch.finfo(dtype).min` at masked positions
+    """
+    if attn_mask.dtype.is_floating_point:
+        mask_value = torch.finfo(attn_mask.dtype).min
+    elif attn_mask.dtype == torch.bool:
+        mask_value = 0
+    else:
+        raise TypeError(f'Unsupported mask type {attn_mask.dtype}')
+    (_, _, query_seq_length, kv_seq_length) = attn_mask.shape
+    causal_mask = torch.ones(attn_mask.shape, dtype=torch.bool, device=attn_mask.device)
+    causal_mask = torch.triu(causal_mask, diagonal=1)
+    causal_mask = causal_mask * (
+        torch.arange(kv_seq_length, device=cache_position.device).view(1, -1) > cache_position.view(-1, 1)
+    ).view(*[1] * (causal_mask.ndim - 2), query_seq_length, kv_seq_length)
+    causal_attn_mask = attn_mask.masked_fill_(causal_mask, mask_value)
+    return causal_attn_mask
+
+
+def update_attn_mask(
+    attn_mask: Optional[torch.Tensor],
+    attn_implementation: str,
+    query_seq_length: int,
+    kv_seq_length: int,
+    cache_position: Optional[torch.Tensor],
+    cache: Optional[transformers.Cache],
+    batch_size: int,
+    causal: bool,
+    dtype: torch.dtype,
+    device: torch.device,
+    output_attentions: bool = False,
+) -> Optional[torch.Tensor]:
+    """
+    Update attn_mask such that it's compatible with the attention implementation.
+    Meant to be used with barrel.train.components.nn.layers.attention.MultiheadAttention and its derivatives
+
+    Args:
+        attn_mask: dtype torch.bool, torch.float32, torch.float16 or torch.bfloat16 and shape one of:
+            - [B, kv_seq_length] (i.e. [B, S])
+            - [B, 1, query_seq_length, kv_seq_length] (i.e. [B, 1, L, S])
+            - [1, 1, query_seq_length, kv_seq_length] (i.e. [L, S])
+            If bool, False values indicate masked positions.
+            If float, must contain only 0.0 and torch.finfo(dtype).min
+            If attn_mask is None, full-bidirectional attention is assumed. The output might be None or
+            a tensor. Refer to the return value documentation
+        attn_implementation: One of [FLASH_ATTN, FLEX_ATTN, SDPA_ATTN, EAGER_ATTN]
+        query_seq_length: The query sequence length (L)
+        kv_seq_length: The key-value sequence length (S)
+        cache_position: dtype torch.int64, shape [query_seq_len]. Used only when causal=True.
+            Contained values are index positions of the query tokens in the sequence. During training,
+            this would usually be torch.arange(query_seq_len), but during generate, this would usually be
+            a tensor sequence with 1 element indicating the position of the token currently being generated.
+            If None, default `cache_positions` are autocomputed from `query_seq_length` and cache size
+        cache: Optional cache. Usually not None when running generate at inference.
+        batch_size: Batch size of the generated attention mask
+        causal: If True, make the attn_mask causal -> all non-causal positions are masked out, regardless
+            of their attn_mask values. When using flash attention or SDPA and `causal == False`, make sure
+            to pass `causal` to the attention operation, in case this function delegates causal masking
+        dtype: dtype of the output attention mask. Must be the dtype of the attn computation
+        device: device of the output attention mask
+        output_attentions: If True, the attention operation is required to output attention weights
+    Returns:
+        - `None` in either of these cases:
+            - `attn_mask` doesn't contain any masked out positions and causal=False
+            - `attn_implementation in [FLASH_ATTN, SDPA_ATTN]` and `attn_mask` doesn't contain any
+                masked out positions. If causal=True, we instead rely on the causal argument to
+                flash attention or `torch.nn.functional.scaled_dot_product_attention`. This happens
+                only if the cache is empty and cache_position is None
+        - `attn_mask` if `attn_implementation == FLASH_ATTN` and `attn_mask` can't be ignored TODO(FLASH)
+        - torch.Tensor of shape [B, 1, query_length, kv_seq_length] (i.e. [B, 1, L, S]) and type `dtype`.
+            Contains zero at unmasked positions and `torch.finfo(dtype).min` at masked positions.
+    """
+    assert attn_implementation in [FLASH_ATTN, FLEX_ATTN, SDPA_ATTN, EAGER_ATTN]
+    assert dtype.is_floating_point, dtype
+    if torch.jit.is_tracing() or torch.jit.is_scripting() or torch.compiler.is_compiling():
+        raise NotImplementedError('Complete correctness not confirmed yet')
+    if isinstance(cache, transformers.StaticCache):
+        if attn_mask is not None and attn_mask.shape[-1] != cache.get_max_cache_shape():
+            raise NotImplementedError('Complete correctness not confirmed yet')
+    full_attn = is_full_attn(attn_mask)
+    past_seen_tokens = cache.get_seq_length() if cache is not None else 0
+    if full_attn and not causal:
+        return None
+    if (
+        full_attn
+        and causal
+        and attn_implementation in [SDPA_ATTN, FLASH_ATTN]
+        and past_seen_tokens == 0
+        and cache_position is None
+    ):
+        return None
+    past_seen_tokens = cache.get_seq_length() if cache is not None else 0
+    static_cache = isinstance(cache, transformers.StaticCache)
+    if static_cache and kv_seq_length < cache.get_max_cache_shape():
+        kv_seq_length = cache.get_max_cache_shape()
+    elif attn_mask is not None:
+        assert kv_seq_length == attn_mask.shape[-1], f'{kv_seq_length}, {attn_mask.shape}'
+    output_mask = make_4d_float_attn_mask(
+        attn_mask=attn_mask,
+        query_seq_length=query_seq_length,
+        kv_seq_length=kv_seq_length,
+        dtype=dtype,
+        device=device,
+        batch_size=batch_size,
+    )
+    if causal:
+        cache_position = (
+            torch.arange(past_seen_tokens, past_seen_tokens + query_seq_length, device=device)
+            if cache_position is None
+            else cache_position
+        )
+        output_mask = make_attn_mask_causal(output_mask, cache_position)
+    if (
+        attn_implementation == SDPA_ATTN
+        and attn_mask is not None
+        and attn_mask.device.type == 'cuda'
+        and not output_attentions
+    ):
+        output_mask = unmask_unattended(output_mask, mask_value=torch.finfo(dtype).min)
+    return output_mask
+
+
+def expand_kv_heads(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    The equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). Convert hidden_states from
+        [batch, num_kv_heads, seqlen, head_dim] -> [batch, num_attention_heads, seqlen, head_dim]
+    """
+    (batch, num_kv_heads, slen, head_dim) = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_kv_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim)
+
+
+def flex_attn_forward(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    attn_mask: Optional[torch.Tensor],
+    num_heads: int,
+    num_kv_heads: int,
+    **kwargs,
+):
+    """
+    Applies FLEX attention on already linearly projected query, key and value.
+    Flex attention should in theory approach flash attention in terms of performance and supports custom
+    2D or 4D attention masks. It uses torch.compile under the hood and requires torch>=2.5.0
+    https://pytorch.org/docs/stable/nn.attention.flex_attention.html
+
+    Args:
+        query_states: Linearly projected query embedding of shape [B, num_heads, L, head_dim]
+        key_states: Linearly projected key embedding of shape [B, num_kv_heads, S, head_dim]
+        value_states: Linearly projected value embedding of shape [B, num_kv_heads, S, head_dim]
+        attn_mask: torch.Tensor of shape [B, 1 | num_heads, L, S] and dtype same as query_states. Contains
+            zeros at unmasked positions and `torch.finfo(attn_mask.dtype).min` at masked positions.
+            If None, no masking is applied.        num_heads: Number of heads for query
+        num_kv_heads: Number of heads for keys and values
+        is_training: True if running in training mode, False otherwise
+        dropout: Dropout probability applied to attention weights
+    Returns:
+        Tuple with entries:
+            - Attention block output: torch.Tensor of shape [B, L, num_heads, head_dim]
+            - None
+    """
+    del kwargs
+    key_states = expand_kv_heads(key_states, num_heads // num_kv_heads)
+    value_states = expand_kv_heads(value_states, num_heads // num_kv_heads)
+    if attn_mask is not None:
+        attn_mask = attn_mask[:, :, :, : key_states.shape[-2]]
+    attn_output = torch.nn.attention.flex_attention.flex_attention(
+        query_states,
+        key_states,
+        value_states,
+        score_mod=(
+            lambda score, batch, head, q_idx, k_idx: score
+            + attn_mask[
+                batch,
+                torch.min(torch.tensor(attn_mask.shape[1] - 1, device=attn_mask.device), head),
+                q_idx,
+                k_idx,
+            ]
+        )
+        if attn_mask is not None
+        else None,
+    )
+    attn_output = attn_output.transpose(1, 2).contiguous()
+    return attn_output, None
+
+
+def sdpa_attn_forward(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    attn_mask: Optional[torch.Tensor],
+    num_heads: int,
+    num_kv_heads: int,
+    dropout: float,
+    is_causal: bool,
+    **kwargs,
+):
+    """
+    Applies SDPA attention on already linearly projected query, key and value via
+    `torch.nn.functional.scaled_dot_product_attention`.
+
+    Args:
+        query_states: Linearly projected query embedding of shape [B, num_heads, L, head_dim]
+        key_states: Linearly projected key embedding of shape [B, num_kv_heads, S, head_dim]
+        value_states: Linearly projected value embedding of shape [B, num_kv_heads, S, head_dim]
+        attn_mask: torch.Tensor of shape [B, 1 | num_heads, L, S] and dtype same as query_states. Contains
+            zeros at unmasked positions and `torch.finfo(attn_mask.dtype).min` at masked positions.
+            If None: no masking is applied if `is_causal` is False and causal mask if `is_causal` is True.
+        dtype torch.bool or same dtype as query/key/value and shape one of:
+        num_heads: Number of heads for query
+        num_kv_heads: Number of heads for keys and values
+        is_training: True if running in training mode, False otherwise
+        dropout: Dropout probability applied to attention weights
+        is_causal: If True, apply additional causal masking whe computing attention
+    Returns:
+        Tuple with entries:
+            - Attention block output: torch.Tensor of shape [B, L, num_heads, head_dim]
+            - None
+    """
+    del kwargs
+    key_states = expand_kv_heads(key_states, num_heads // num_kv_heads)
+    value_states = expand_kv_heads(value_states, num_heads // num_kv_heads)
+    if attn_mask is not None:
+        attn_mask = attn_mask[:, :, :, : key_states.shape[-2]]
+    attn_output = torch.nn.functional.scaled_dot_product_attention(
+        query_states, key_states, value_states, attn_mask=attn_mask, dropout_p=dropout, is_causal=is_causal
+    )
+    attn_output = attn_output.transpose(1, 2).contiguous()
+    return attn_output, None
+
+
+def eager_attn_forward(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    attn_mask: Optional[torch.Tensor],
+    num_heads: int,
+    num_kv_heads: int,
+    dropout: float,
+    **kwargs,
+):
+    """
+    Applies EAGER attention on already linearly projected query, key and value.
+
+    Args:
+        query_states: Linearly projected query embedding of shape [B, num_heads, L, head_dim]
+        key_states: Linearly projected key embedding of shape [B, num_kv_heads, S, head_dim]
+        value_states: Linearly projected value embedding of shape [B, num_kv_heads, S, head_dim]
+        attn_mask: torch.Tensor of shape [B, 1 | num_heads, L, S] and dtype same as query_states. Contains
+            zeros at unmasked positions and `torch.finfo(attn_mask.dtype).min` at masked positions.
+            If None, no masking is applied.
+        num_heads: Number of heads for query
+        num_kv_heads: Number of heads for keys and values
+        is_training: True if running in training mode, False otherwise
+        dropout: Dropout probability applied to attention weights
+        is_causal: If True, apply additional causal masking whe computing attention
+    Returns:
+        Tuple with entries:
+            - Attention block output: torch.Tensor of shape [B, L, num_heads, head_dim]
+            - Attention weights: torch.Tensor of shape [B, num_heads, L, S]
+    """
+    del kwargs
+    head_dim = key_states.shape[-1]
+    key_states = expand_kv_heads(key_states, num_heads // num_kv_heads)
+    value_states = expand_kv_heads(value_states, num_heads // num_kv_heads)
+    attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(head_dim)
+    if attn_mask is not None:
+        attn_mask = attn_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + attn_mask
+    attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
+        query_states.dtype
+    )
+    attn_weights = torch.nn.functional.dropout(attn_weights, p=dropout)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+    return attn_output, attn_weights
+
+
+ATTN_TYPES = {
+    EAGER_ATTN: eager_attn_forward,
+    SDPA_ATTN: sdpa_attn_forward,
+    FLEX_ATTN: flex_attn_forward,
+    FLASH_ATTN: flash_attn_2_forward,
+}
+
+
+class MultiheadAttention(torch.nn.Module):
+    """
+    Multi-headed attention from 'Attention Is All You Need' paper
+
+    Different implementation from torch.nn.MultiheadAttention to support:
+        - Easy switch between EAGER_ATTN, SDPA_ATTN and FLASH_ATTN
+        - Number of key-value heads different from query heads
+        - Key-value cache during forward, in the same way as transformers. Useful for generation or
+            cross-attention to projected keys and values
+        - Ability to apply positional encodings to key and value after input linear projection
+        - Different linear projection output size
+
+    Adapted from transformers.models.llama.modeling_llama.LlamaAttention
+    """
+
+    def __init__(
+        self,
+        attn_implementation: str,
+        in_features: int,
+        num_heads: int,
+        head_dim: Optional[int] = None,
+        out_features: Optional[int] = None,
+        key_features: Optional[int] = None,
+        value_features: Optional[int] = None,
+        num_kv_heads: Optional[int] = None,
+        bias: bool = False,
+        dropout: float = 0.0,
+        cache_layer: Optional[int] = None,
+        query_position_embed: Optional[Callable[[torch.Tensor, Optional[torch.Tensor]], torch.Tensor]] = None,
+        key_position_embed: Optional[Callable[[torch.Tensor, Optional[torch.Tensor]], torch.Tensor]] = None,
+    ):
+        """
+        Args:
+            in_features: Input dimension for query linear projection.
+            num_heads: Number of heads for query
+            head_dim: Head dimension. If None, defaults to `in_features // num_heads`
+            out_features: Output dimension for the output linear layer. If None, defaults to `in_features`
+            key_features: Input dimension for key linear projection. If None, defaults to `in_features`
+            value_features: Input dimension for value linear projection. If None, defaults to `in_features`
+            num_kv_heads: Number of heads for keys and values. If None, defaults to `num_heads`
+            cache_layer: Index of the layer in the cache. Needed only when `cache` is passed to
+                the `forward()` call, usually during generation or when the projected keys and values need
+                to be cached during training. Can be omitted when `cache_layer` is passed to `forward`
+            position_embed: Callable that takes as input linearly projected query and key and a tuple of
+                positional embeddings and returns query and key with positional embeddings applied. Note
+                these embeddings are applied after linear projection. If you want to apply embeddings before
+                the linear projection, do so before calling the forward method and use the default value
+                for `position_embed`, which is a simple pass-through. Note you can also pass torch.nn.Module
+            key_position_embed: Callable that takes as input linearly projected key and optional positional
+                index in the sequence and returns key with positional embeddings applied.
+                positional embeddings and returns query and key with positional embeddings applied. Note
+                these embeddings are applied after linear projection. If you want to apply embeddings before
+                the linear projection, do so before calling the forward method and use the default value
+                for `position_embed`, which is a simple pass-through. Note you can also pass torch.nn.Module
+        """
+        super().__init__()
+        assert attn_implementation in ATTN_TYPES, attn_implementation
+        self.attn_implementation = attn_implementation
+        self.attn_forward = ATTN_TYPES[attn_implementation]
+        self.in_features = in_features
+        self.key_features = key_features or in_features
+        self.value_features = value_features or in_features
+        self.bias = bias
+        self.out_features = out_features or in_features
+        self.num_heads = num_heads
+        self.head_dim = head_dim or in_features // num_heads
+        self.num_kv_heads = num_kv_heads or num_heads
+        self.dropout = dropout
+        self.query_position_embed = query_position_embed
+        self.key_position_embed = key_position_embed
+        self.cache_layer = cache_layer
+        self.q_proj = torch.nn.Linear(self.in_features, self.num_heads * self.head_dim, bias=self.bias)
+        self.k_proj = torch.nn.Linear(self.key_features, self.num_kv_heads * self.head_dim, bias=self.bias)
+        self.v_proj = torch.nn.Linear(self.value_features, self.num_kv_heads * self.head_dim, bias=self.bias)
+        self.o_proj = torch.nn.Linear(self.num_heads * self.head_dim, self.out_features, bias=self.bias)
+        if self.attn_implementation == FLEX_ATTN:
+            assert self.dropout == 0.0, "FLEX attention doesn't support dropout"
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        is_causal: bool = False,
+        query_position_indices: Optional[torch.Tensor] = None,
+        key_position_indices: Optional[torch.Tensor] = None,
+        cache: Optional[transformers.Cache] = None,
+        cache_layer: Optional[int] = None,
+        output_attentions: bool = False,
+        cache_kwargs: Dict[str, Any] = {},
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """
+        Args:
+            query: Query embedding of shape [B, L, in_features]
+            key: Key embedding of shape [B, S, key_features]
+            value: Value embedding of shape [B, S, value_features]
+            attn_mask: dtype torch.bool or same dtype as query/key/value and shape one of:
+                - [B, S]
+                - [B | 1, 1 | num_heads, L, S]
+                If bool, False values indicate masked positions (opposite of torch.nn.MultiheadAttention)
+                If float, must contain only 0.0 and torch.finfo(dtype).min
+                If attn_mask is None, full-bidirectional attention or causal attention is used depdening
+                on the value of `is_causal`.
+                If FLASH_ATTN is used as `attn_implementation`, only bool attn_mask of shape [B, S]
+                or None is supported.
+            is_causal: If True, apply additional causal masking to `attn_mask`
+            query_position_indices: torch.Tensor of shape [1 | B, L] containing the indices of the `query`
+                tokens within the entire sequence. Passed through to query_position_embed. If None and `cache`
+                is not None, indices are autogenerated [0, 1, ..., L] and offset by `cache_size`
+            key_position_indices: Same as `query_position_indices`, but applied to key
+            cache: transformers.Cache containing cached key-value pairs. The linearly projected
+                `key` and `value` passed to this function get added to the cache and concatenated after the
+                key-value pairs in the cache and then attention is computed on the concatenated sequence.
+                This is most commonly used at inference when generating auto-regressively or when one needs
+                to cross attend to the keys and values outside this module forward pass.
+            cache_layer: Index of the layer in the cache. Needed only when `cache` is passed to
+                the `forward()` call, usually during generation or when the projected keys and values need
+                to be cached during training. Can be omitted when `cache_layer` was passed to `__init__`
+            output_attentions: If True, output also the attention weights. Otherwise output None.
+                Note that only the eager implementation of MultiheadAttention supports this.
+            cache_kwargs: kwargs directly passed to `cache.update()`
+        Returns:
+            Tuple with entries:
+                - Attention block output: torch.Tensor of shape [B, L, out_features]
+                - Optional attention weights if `output_attentions=True`, shape [B, num_heads, L, S]
+        """
+        if self.attn_implementation != EAGER_ATTN:
+            assert (
+                output_attentions is False
+            ), f"{self.attn_implementation} doesn't support output_attentions=True"
+        batch_size = query.shape[0]
+        query_states = self.q_proj(query)
+        key_states = self.k_proj(key)
+        value_states = self.v_proj(value)
+        query_states = query_states.view(
+            batch_size, query_states.shape[1], self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        key_states = key_states.view(
+            batch_size, key_states.shape[1], self.num_kv_heads, self.head_dim
+        ).transpose(1, 2)
+        value_states = value_states.view(
+            batch_size, value_states.shape[1], self.num_kv_heads, self.head_dim
+        ).transpose(1, 2)
+        (query_states, key_states) = self._maybe_apply_positional_embeddings(
+            query_states=query_states,
+            key_states=key_states,
+            query_position_indices=query_position_indices,
+            key_position_indices=key_position_indices,
+            cache=cache,
+        )
+        (key_states, value_states) = self._maybe_update_cache(
+            key_states, value_states, cache_layer=cache_layer, cache=cache, cache_kwargs=cache_kwargs
+        )
+        attn_mask = update_attn_mask(
+            attn_mask,
+            attn_implementation=self.attn_implementation,
+            query_seq_length=query_states.shape[2],
+            kv_seq_length=value_states.shape[2],
+            cache_position=query_position_indices,
+            cache=cache,
+            batch_size=batch_size,
+            causal=is_causal,
+            dtype=query_states.dtype,
+            device=query_states.device,
+            output_attentions=output_attentions,
+        )
+        dropout = self.dropout if self.training else 0.0
+        (attn_output, attn_weights) = self.attn_forward(
+            query_states=query_states,
+            key_states=key_states,
+            value_states=value_states,
+            attn_mask=attn_mask,
+            num_heads=self.num_heads,
+            num_kv_heads=self.num_kv_heads,
+            is_causal=is_causal,
+            dropout=dropout,
+        )
+        shape = (batch_size, query.shape[1], self.num_heads, self.head_dim)
+        assert attn_output.shape == shape, f'{attn_output.shape} != {shape}'
+        attn_output = attn_output.reshape(batch_size, query.shape[1], self.num_heads * self.head_dim)
+        attn_output = self.o_proj(attn_output)
+        if not output_attentions:
+            attn_weights = None
+        return attn_output, attn_weights
+
+    def _maybe_apply_positional_embeddings(
+        self,
+        query_states: torch.Tensor,
+        key_states: torch.Tensor,
+        query_position_indices: Optional[torch.Tensor],
+        key_position_indices: Optional[torch.Tensor],
+        cache: Optional[transformers.Cache],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        device = query_states.device
+        if self.query_position_embed is not None:
+            if query_position_indices is None and cache is not None:
+                query_position_indices = (
+                    torch.arange(query_states.shape[-2], dtype=torch.int64, device=device).view(1, -1)
+                    + cache.get_seq_length()
+                )
+            query_states = self.query_position_embed(query_states, position_indices=query_position_indices)
+        if self.key_position_embed is not None:
+            if key_position_indices is None and cache is not None:
+                key_position_indices = (
+                    torch.arange(key_states.shape[-2], dtype=torch.int64, device=device).view(1, -1)
+                    + cache.get_seq_length()
+                )
+            key_states = self.key_position_embed(key_states, position_indices=key_position_indices)
+        return query_states, key_states
+
+    def _maybe_update_cache(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        cache_layer: Optional[int],
+        cache: Optional[transformers.Cache],
+        cache_kwargs: Dict[str, Any],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if cache is not None:
+            if cache_layer is None and self.cache_layer is None:
+                raise RuntimeError('When cache != None, cache_layer has to be set')
+            cache_layer = cache_layer if cache_layer is not None else self.cache_layer
+            (key_states, value_states) = cache.update(key_states, value_states, cache_layer, cache_kwargs)
+        return key_states, value_states
+
+
+def make_activation(activation: str | Type[torch.nn.Module], **kwargs) -> torch.nn.Module:
+    if isinstance(activation, str):
+        TorchActivation: Type[torch.nn.Module] = getattr(torch.nn, activation)
+    else:
+        TorchActivation: Type[torch.nn.Module] = activation
+    assert issubclass(TorchActivation, torch.nn.Module), TorchActivation
+    return TorchActivation(**kwargs)
+
+
+class PiZeroMLP(torch.nn.Module):
+    def __init__(
+        self, feature_size: int, hidden_size: int, activation: str, activation_kwargs: Dict[str, Any] = {}
+    ):
+        super().__init__()
+        self.gate_proj = torch.nn.Linear(feature_size, hidden_size, bias=False)
+        self.up_proj = torch.nn.Linear(feature_size, hidden_size, bias=False)
+        self.down_proj = torch.nn.Linear(hidden_size, feature_size, bias=False)
+        self.activation = make_activation(activation, **activation_kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x))
+
+
+class PiZeroFlowMatchingDecoderBlock(ConfigurableModule[PiZeroFlowMatchingDecoderBlockConfig]):
+    def __init__(self, config: PiZeroFlowMatchingDecoderBlockConfig, **attn_kwargs):
+        super().__init__(config)
+        self.norm_in = GemmaRMSNorm(self.config.feature_size, eps=1e-06)
+        self.self_attn = MultiheadAttention(
+            attn_implementation=self.config.attn_implementation,
+            in_features=self.config.feature_size,
+            num_heads=self.config.num_heads,
+            head_dim=self.config.head_dim,
+            num_kv_heads=self.config.num_kv_heads,
+            **attn_kwargs,
+        )
+        self.mlp = PiZeroMLP(
+            feature_size=self.config.feature_size,
+            hidden_size=self.config.hidden_size,
+            activation=self.config.activation,
+            activation_kwargs=self.config.activation_kwargs,
+        )
+        self.norm_out = GemmaRMSNorm(self.config.feature_size, eps=1e-06)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        attn_mask: torch.Tensor,
+        cache: transformers.Cache,
+        attn_kwargs: Dict[str, Any],
+    ) -> torch.Tensor:
+        """
+        Args:
+            query: torch.Tensor of shape [B, L, token_size]. The query seqence in the order:
+                [noised query tokens, condition token, robot state tokens]
+            timestep: torch.Tensor of shape [B, 1, token_size]. Timestep token
+            attn_mask: torch.Tensor of shape [B, 1, L, L+S] and dtype torch.bool, where S is the VLM
+                sequence length
+            cache: Cache that contains only the VLM tokens during training and VLM + past query tokens
+                during generation
+            num_noised_tokens: Number of noised tokens in `query`
+            num_condition_tokens: Number of condition tokens in `query`
+        Returns:
+            torch.Tensor of same shape as query [B, L, token_size]
+        """
+        residual = x = query
+        x = self.norm_in(x)
+        (x, _) = self.self_attn(
+            query=x, key=x, value=x, attn_mask=attn_mask, is_causal=False, cache=cache, **attn_kwargs
+        )
+        x = residual + x
+        residual = x
+        x = self.norm_out(x)
+        x = self.mlp(x)
+        x = residual + x
+        return x
+
+
+class PiZeroFlowMatchingDecoder(ConfigurableModule[PiZeroFlowMatchingDecoderConfig]):
+    """PiZero Flow Matching control decoder"""
+
+    def __init__(self, config: PiZeroFlowMatchingDecoderConfig):
+        super().__init__(config)
+        query_position_embed = RotaryPositionalEncoding(config=self.config.block_config.position_embed_config)
+        key_position_embed = RotaryPositionalEncoding(config=self.config.block_config.position_embed_config)
+        self.blocks = torch.nn.ModuleList(
+            [
+                PiZeroFlowMatchingDecoderBlock(
+                    self.config.block_config,
+                    query_position_embed=query_position_embed,
+                    key_position_embed=key_position_embed,
+                    cache_layer=i,
+                )
+                for i in range(self.config.num_blocks)
+            ]
+        )
+        self.norm = GemmaRMSNorm(self.config.block_config.feature_size, eps=1e-06)
+
+    def forward(
+        self,
+        control_tokens: torch.Tensor,
+        robot_state_tokens: torch.Tensor,
+        llm_kv_tokens: List[Tuple[torch.Tensor, torch.Tensor]],
+        attn_mask: Optional[torch.Tensor],
+        cache: Optional[transformers.StaticCache] = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            control_tokens: torch.Tensor of shape [B, N, token_size], contains sequence of controls
+            robot_state_tokens: torch.Tensor of shape [B, num_state_tokens, token_size]
+            llm_kv_tokens: List of linearly projected key-value pairs from LLM, right before attention
+                operation. Each tensor is of the shape [B, num_kv_heads, S, head_dim]
+            attn_mask: One of
+                - shape [B, S], dtype torch.bool -> padding attention mask for LLM tokens
+                - shape [B, 1, L, S], dtype torch.bool -> full attention mask for LLM tokens
+            cache:
+                - When None, we are either in training mode or generation mode without cache. In the latter
+                case, this means we don't cache the robot state key value pairs, but compute them every time
+                - When provided, we are in generation mode with cache. The cache could be empty (step zero)
+                or contain both the VLM key value pairs and past robot state key value pairs (non-zero step).
+                Furthermore, the cache state is updated and preserved across generation steps.
+        Returns:
+            torch.Tensor, shape [B, N, token_size]
+        """
+        assert (
+            len(llm_kv_tokens) == self.config.num_blocks
+        ), f'{len(llm_kv_tokens)} != {self.config.num_blocks}'
+        cache_is_empty = cache.get_seq_length() == 0 if cache is not None else True
+        vlm_seq_len = attn_mask.shape[-1]
+        device = attn_mask.device
+        if cache is None:
+            cache = transformers.DynamicCache()
+        if cache_is_empty:
+            position_indices = torch.arange(vlm_seq_len, dtype=torch.int64, device=device)
+            for block_index, kv_tokens in enumerate(llm_kv_tokens):
+                (key_states, value_states) = kv_tokens
+                cache.update(
+                    key_states, value_states, block_index, cache_kwargs={'cache_position': position_indices}
+                )
+        num_control_tokens = control_tokens.shape[1]
+        num_robot_state_tokens = robot_state_tokens.shape[1]
+        attn_mask = self._build_attn_mask(
+            num_control_tokens=num_control_tokens,
+            num_robot_state_tokens=num_robot_state_tokens,
+            attn_mask=attn_mask,
+        )
+        if cache_is_empty:
+            tokens = torch.cat([robot_state_tokens, control_tokens], axis=1)
+            query_position_indices = key_position_indices = vlm_seq_len + torch.arange(
+                tokens.shape[1], dtype=torch.int64, device=device
+            ).view(1, -1)
+        else:
+            tokens = control_tokens
+            attn_mask = attn_mask[:, :, -control_tokens.shape[1] :]
+            query_position_indices = key_position_indices = (
+                vlm_seq_len
+                + num_robot_state_tokens
+                + torch.arange(tokens.shape[1], dtype=torch.int64, device=device).view(1, -1)
+            )
+        for block in self.blocks:
+            tokens = block(
+                query=tokens,
+                attn_mask=attn_mask,
+                cache=cache,
+                attn_kwargs={
+                    'query_position_indices': query_position_indices,
+                    'key_position_indices': key_position_indices,
+                    'cache_kwargs': {'cache_position': key_position_indices.view(-1)},
+                },
+            )
+        if cache_is_empty:
+            (_, control_tokens) = torch.split(tokens, [num_robot_state_tokens, num_control_tokens], dim=1)
+        else:
+            control_tokens = tokens
+        control_tokens = self.norm(control_tokens)
+        return control_tokens
+
+    @torch.no_grad()
+    def _build_attn_mask(
+        self, num_control_tokens: int, num_robot_state_tokens: int, attn_mask: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Expand `attn_mask` (which is effectively a padding mask) to 4D such that:
+            - robot state tokens and control tokens can't attend to padding tokens
+            - robot state tokens can't attend to control tokens
+        Note: We can't keep the mask in 2D as it doesn't allow masking of padding tokens from the
+        VLM sequence. Furthermore, in a 2D mask you can't disable attention from robot state tokens
+        to control tokens
+        """
+        assert attn_mask.dtype == torch.bool, attn_mask.dtype
+        assert attn_mask.ndim in [2, 4], attn_mask.shape
+        device = attn_mask.device
+        batch_size = attn_mask.shape[0]
+        query_seq_len = num_robot_state_tokens + num_control_tokens
+        vlm_seq_len = attn_mask.shape[-1]
+        kv_seq_len = query_seq_len + vlm_seq_len
+        cross_attn_mask = torch.ones(
+            [batch_size, 1, query_seq_len, kv_seq_len], dtype=torch.bool, device=device
+        )
+        if attn_mask.ndim == 2:
+            attn_mask = attn_mask.view(batch_size, 1, 1, vlm_seq_len)
+        else:
+            attn_mask = torch.any(attn_mask, dim=-2, keepdims=True)
+        cross_attn_mask[..., :vlm_seq_len] = attn_mask
+        robot_state_query_indices = torch.arange(
+            num_robot_state_tokens, dtype=torch.int64, device=device
+        ).view(-1, 1)
+        control_key_indices = (
+            torch.arange(num_control_tokens, dtype=torch.int64, device=device).view(-1, 1)
+            + vlm_seq_len
+            + num_robot_state_tokens
+        )
+        cross_attn_mask[:, :, robot_state_query_indices, control_key_indices] = 0
+        return cross_attn_mask
+
+    @property
+    def fsdp_wrap_modules(self) -> Dict[torch.nn.Module, Dict[str, Any]]:
+        return {
+            **{module: {} for module in self.modules() if isinstance(module, type(self.blocks[0]))},
+            self.norm: {},
+        }
+
+
+class VLMInput(Protocol):
+    input_ids: torch.Tensor
+    attn_mask: torch.Tensor
+    images: Dict[str, torch.Tensor]
+    multimodal_indices: torch.Tensor
+    unimodal_indices: torch.Tensor
+
+    @property
+    def inputs_embeds(self) -> Optional[torch.Tensor]:
+        return None
+
+    @property
+    def past_key_values(self) -> Optional[List[torch.Tensor]]:
+        return None
+
+
+VLMConfigT = TypeVar('VLMConfigT', bound=VLMConfig)
+
+
+class VLM(ConfigurableModule[VLMConfigT], Template[VLMConfigT]):
+    """
+    Abstract class for arbitrary Vision-Language Models
+
+    Explicitly don't inherit from `transformers.PretrainedModel` or any other `transformers` subclasses.
+    Instead, keep 'compatible' APIs such that the underlying `generate` utilities of `transformers` can
+    be used via composition by classes that have instances of this class as an attribute
+    """
+
+    @property
+    @abstractmethod
+    def fsdp_wrap_modules(self) -> Dict[torch.nn.Module, Dict[str, Any]]:
+        ...
+
+    @abstractmethod
+    def forward(
+        self,
+        inputs: VLMInput,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs,
+    ) -> VLMOutput:
+        ...
+
+
+def qwen3_vl_mixed_modality_forward(
+    self,
+    input_ids: torch.LongTensor = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[transformers.cache_utils.Cache] = None,
+    inputs_embeds: Optional[torch.FloatTensor] = None,
+    pixel_values: Optional[torch.Tensor] = None,
+    pixel_values_videos: Optional[torch.FloatTensor] = None,
+    image_grid_thw: Optional[torch.LongTensor] = None,
+    video_grid_thw: Optional[torch.LongTensor] = None,
+    cache_position: Optional[torch.LongTensor] = None,
+    second_per_grid_ts: Optional[torch.Tensor] = None,
+    **kwargs,
+) -> Union[tuple, transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLModelOutputWithPast]:
+    """
+    Adapted from:
+    https://github.com/2U1/Qwen-VL-Series-Finetune/blob/512f424e74f94755d774b6e3786457750677048b/src/train/monkey_patch_forward.py#L173
+    """
+    del second_per_grid_ts
+    if (input_ids is None) ^ (inputs_embeds is not None):
+        raise ValueError('You must specify exactly one of input_ids or inputs_embeds')
+    if inputs_embeds is None:
+        inputs_embeds = self.get_input_embeddings()(input_ids)
+    image_mask = None
+    video_mask = None
+    if pixel_values is None and pixel_values_videos is None:
+        dummy_pixel = torch.zeros(1024, 1536).to(self.visual.device)
+        dummy_grid = torch.tensor([[1, 32, 32]]).to(self.visual.device)
+        (image_embeds, dummy_deepstack) = self.get_image_features(dummy_pixel, dummy_grid)
+        image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+        inputs_embeds += image_embeds.mean() * 0
+    if pixel_values is not None:
+        (image_embeds, deepstack_image_embeds) = self.get_image_features(pixel_values, image_grid_thw)
+        image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+        (image_mask, _) = self.get_placeholder_mask(
+            input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
+        )
+        inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+    if pixel_values_videos is not None:
+        (video_embeds, deepstack_video_embeds) = self.get_video_features(pixel_values_videos, video_grid_thw)
+        video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+        (_, video_mask) = self.get_placeholder_mask(
+            input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds
+        )
+        inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+    visual_pos_masks = None
+    deepstack_visual_embeds = None
+    if image_mask is not None and video_mask is not None:
+        image_mask = image_mask[..., 0]
+        video_mask = video_mask[..., 0]
+        visual_pos_masks = image_mask | video_mask
+        deepstack_visual_embeds = []
+        image_mask_joint = image_mask[visual_pos_masks]
+        video_mask_joint = video_mask[visual_pos_masks]
+        for img_embed, vid_embed in zip(deepstack_image_embeds, deepstack_video_embeds, strict=False):
+            embed_joint = img_embed.new_zeros(visual_pos_masks.sum(), img_embed.shape[-1]).to(
+                img_embed.device
+            )
+            embed_joint[image_mask_joint, :] = img_embed
+            embed_joint[video_mask_joint, :] = vid_embed
+            deepstack_visual_embeds.append(embed_joint)
+    elif image_mask is not None:
+        image_mask = image_mask[..., 0]
+        visual_pos_masks = image_mask
+        deepstack_visual_embeds = deepstack_image_embeds
+    elif video_mask is not None:
+        video_mask = video_mask[..., 0]
+        visual_pos_masks = video_mask
+        deepstack_visual_embeds = deepstack_video_embeds
+    if visual_pos_masks is None:
+        (B, S, _) = inputs_embeds.shape
+        visual_pos_masks = torch.zeros((B, S), dtype=torch.bool, device=inputs_embeds.device)
+        deepstack_visual_embeds = [t.narrow(0, 0, 0) for t in dummy_deepstack]
+    if position_ids is None:
+        attention_mask_tensor = (
+            attention_mask if not isinstance(attention_mask, dict) else attention_mask['full_attention']
+        )
+        if attention_mask_tensor is not None and attention_mask_tensor.ndim == 4:
+            attention_mask_tensor = torch.diagonal(attention_mask_tensor[:, 0], dim1=1, dim2=2)
+            if attention_mask_tensor.dtype.is_floating_point:
+                attention_mask_tensor = attention_mask_tensor / torch.finfo(attention_mask_tensor.dtype).min
+                attention_mask_tensor = (1.0 - attention_mask_tensor).int()
+        prefill_compiled_stage = transformers.utils.is_torchdynamo_compiling() and (
+            input_ids is not None
+            and input_ids.shape[1] != 1
+            or inputs_embeds is not None
+            and inputs_embeds.shape[1] != 1
+        )
+        prefill_noncompiled_stage = not transformers.utils.is_torchdynamo_compiling() and (
+            cache_position is not None
+            and cache_position[0] == 0
+            or past_key_values is None
+            or past_key_values.get_seq_length() == 0
+        )
+        if prefill_compiled_stage or prefill_noncompiled_stage or self.rope_deltas is None:
+            (position_ids, rope_deltas) = self.get_rope_index(
+                input_ids, image_grid_thw, video_grid_thw, attention_mask=attention_mask_tensor
+            )
+            self.rope_deltas = rope_deltas
+        else:
+            (batch_size, seq_length, _) = inputs_embeds.shape
+            delta = (
+                (cache_position[0] + self.rope_deltas).to(inputs_embeds.device)
+                if cache_position is not None
+                else 0
+            )
+            position_ids = torch.arange(seq_length, device=inputs_embeds.device)
+            position_ids = position_ids.view(1, -1).expand(batch_size, -1)
+            if cache_position is not None:
+                delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0)
+            position_ids = position_ids.add(delta)
+            position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+    outputs = self.language_model(
+        input_ids=None,
+        position_ids=position_ids,
+        attention_mask=attention_mask,
+        past_key_values=past_key_values,
+        inputs_embeds=inputs_embeds,
+        cache_position=cache_position,
+        visual_pos_masks=visual_pos_masks,
+        deepstack_visual_embeds=deepstack_visual_embeds,
+        **kwargs,
+    )
+    return transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLModelOutputWithPast(
+        last_hidden_state=outputs.last_hidden_state,
+        past_key_values=outputs.past_key_values,
+        rope_deltas=self.rope_deltas,
+    )
+
+
+def replace_qwen3_vl_with_mixed_modality_forward():
+    """
+    Adapted from: https://github.com/2U1/Qwen-VL-Series-Finetune/blob/512f424e74f94755d774b6e3786457750677048b/src/train/monkey_patch_forward.py#L21
+    """
+    transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLModel.forward = qwen3_vl_mixed_modality_forward
+
+
+class Qwen3VL(VLM[Qwen3VLConfig]):
+    def __init__(self, config: Qwen3VLConfig):
+        super().__init__(config)
+        if self.config.mixed_modality_forward:
+            replace_qwen3_vl_with_mixed_modality_forward()
+        self.model = (
+            transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLForConditionalGeneration.from_pretrained(
+                config.model_id, attn_implementation=config.attn_implementation
+            )
+        )
+        if not self.config.lm_head:
+            self.model.lm_head = torch.nn.Identity()
+        hf_processor = transformers.AutoProcessor.from_pretrained(config.model_id)
+        self.processor = Qwen3VLProcessor(config=self.config.processor_config, hf_processor=hf_processor)
+        self.model.train()
+
+    def _flatten_and_unpad_pixel_values(
+        self, pixel_values: torch.Tensor, grid_thw: torch.Tensor
+    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+        (B, N, C) = pixel_values.shape
+        pixel_values = pixel_values.view(B * N, C)
+        patch_mask = (pixel_values < -100.0).any(dim=-1)
+        pixel_values = pixel_values[~patch_mask]
+        grid_thw = grid_thw.reshape(-1, 3)
+        grid_mask = (grid_thw < 0).any(dim=-1)
+        grid_thw = grid_thw[~grid_mask]
+        if pixel_values.shape[0] == 0 or grid_thw.shape[0] == 0:
+            return None, None
+        assert (
+            pixel_values.shape[0] == (grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).sum().item()
+        ), "Number of patches doesn't match the grid dimensions."
+        return pixel_values, grid_thw
+
+    def _prepare_vision_inputs(
+        self, images: Dict[str, torch.Tensor]
+    ) -> Tuple[
+        Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]
+    ]:
+        pixel_values = images.get('pixel_values', None)
+        pixel_values_videos = images.get('pixel_values_videos', None)
+        image_grid_thw = images.get('image_grid_thw', None)
+        video_grid_thw = images.get('video_grid_thw', None)
+        if pixel_values is None and pixel_values_videos is None:
+            raise ValueError(
+                "Either 'pixel_values' or 'pixel_values_videos' must be provided in images dict."
+            )
+        if pixel_values is not None:
+            if image_grid_thw is None:
+                raise ValueError(
+                    "'image_grid_thw' must be provided in images dict when 'pixel_values' is provided."
+                )
+            (pixel_values, image_grid_thw) = self._flatten_and_unpad_pixel_values(
+                pixel_values, image_grid_thw
+            )
+        if pixel_values_videos is not None:
+            if video_grid_thw is None:
+                raise ValueError(
+                    "'video_grid_thw' must be provided in images dict when 'pixel_values_videos' is provided."
+                )
+            raise NotImplementedError('Video input not yet supported for Qwen3-VL.')
+        return pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw
+
+    def forward(
+        self,
+        inputs: VLMInput,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        **kwargs,
+    ) -> VLMOutput:
+        del kwargs
+        (pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw) = self._prepare_vision_inputs(
+            inputs.images
+        )
+        cache = transformers.DynamicCache()
+        model_input_args = dict(
+            input_ids=inputs.input_ids,
+            pixel_values=pixel_values,
+            pixel_values_videos=pixel_values_videos,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            attention_mask=inputs.attn_mask,
+            use_cache=use_cache,
+            past_key_values=cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+        )
+        if self.config.lm_head:
+            llm_output: transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLCausalLMOutputWithPast = (
+                self.model(**model_input_args)
+            )
+        else:
+            llm_output: transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLModelOutputWithPast = (
+                self.model.model(**model_input_args)
+            )
+        image_mask = inputs.input_ids == self.processor.hf_processor.image_token_id
+        text_mask = (inputs.input_ids != self.processor.ignore_index) & ~image_mask
+        output = VLMOutput(
+            llm_output=LLMOutput.from_transformers(
+                input_ids=inputs.input_ids, llm_output=llm_output, text_mask=text_mask, image_mask=image_mask
+            ),
+            vit_tokens=None,
+            attn_mask=inputs.attn_mask,
+        )
+        return output
+
+    @property
+    def fsdp_wrap_modules(self) -> Dict[torch.nn.Module, Dict[str, Any]]:
+        transformer_modules = {
+            module: {}
+            for module in self.modules()
+            if isinstance(
+                module,
+                (
+                    transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionBlock,
+                    transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLTextDecoderLayer,
+                    transformers.models.qwen3_vl.modeling_qwen3_vl.Qwen3VLVisionPatchMerger,
+                ),
+            )
+            or module in (self.model.language_model.embed_tokens, self.model.language_model.norm)
+        }
+        if not self.config.lm_head:
+            transformer_modules[self.model.language_model.layers[-1]] = {'reshard_after_forward': False}
+        return transformer_modules
+
+    @torch.inference_mode()
+    def generate(self, inputs: VLMInput, do_sample: bool = False, max_new_tokens: int = 512) -> VLMOutput:
+        assert self.config.lm_head, 'Generation is only supported when lm_head is present.'
+        (pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw) = self._prepare_vision_inputs(
+            inputs.images
+        )
+        vlm_output: transformers.generation.utils.GenerateBeamDecoderOnlyOutput = self.model.generate(
+            input_ids=inputs.input_ids,
+            pixel_values=pixel_values,
+            pixel_values_videos=pixel_values_videos,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            attention_mask=inputs.attn_mask,
+            do_sample=do_sample,
+            max_new_tokens=max_new_tokens,
+            return_dict_in_generate=True,
+        )
+        image_mask = inputs.input_ids == self.processor.hf_processor.image_token_id
+        text_mask = (inputs.input_ids != self.processor.ignore_index) & ~image_mask
+        output = VLMOutput(
+            llm_output=LLMOutput.make_empty().replace(
+                input_ids=inputs.input_ids,
+                output_ids=vlm_output.sequences,
+                past_key_values=list(vlm_output.past_key_values),
+                text_mask=text_mask,
+                image_mask=image_mask,
+            ),
+            vit_tokens=None,
+            attn_mask=inputs.attn_mask,
+        )
+        return output
+
+
+def integrate_unitquat(
+    qt: torch.Tensor,
+    dq_dt: torch.Tensor,
+    dt: float | torch.Tensor,
+    body_frame: bool = True,
+    half_cover: bool = True,
+) -> torch.Tensor:
+    """
+    Integrate a unit quaternion `qt` by the derivative `dq_dt` over the time interval `dt`.
+    Args:
+        qt: Unit quaternion, shape [..., 4]
+        dq_dt: Derivative of the unit quaternion, shape [..., 4]
+        dt: Time interval to integrate over, scalar or a tensor of shape () or [..., 1]
+        half_cover: If True, the result is guaranteed to lie in the half space
+        body_frame: If True, the integration is done in the body frame (post-multiply),
+            otherwise in the inertial frame (pre-multiply).
+    Returns:
+        Integrated unit quaternion, shape [..., 4]
+    """
+    assert qt.shape == dq_dt.shape, f'{qt.shape} != {dq_dt.shape}'
+    assert is_quaternion(qt), f'{qt.shape} not a quaternion'
+    if isinstance(dt, torch.Tensor):
+        assert dt.ndim in (0, qt.ndim), f'dt.ndim = {dt.ndim} | {qt.ndim}'
+    if body_frame:
+        omega_q = 2.0 * roma.quat_product(roma.quat_conjugation(qt), dq_dt)
+    else:
+        omega_q = 2.0 * roma.quat_product(dq_dt, roma.quat_conjugation(qt))
+    omega = omega_q[..., :-1]
+    dq = roma.rotvec_to_unitquat(omega * dt)
+    if body_frame:
+        qt = roma.quat_product(qt, dq)
+    else:
+        qt = roma.quat_product(dq, qt)
+    if half_cover:
+        qt = quaternion_half_cover(qt)
+    return qt
+
+
+def skew_symmetric_to_rotvec(skew_symmetric: torch.Tensor) -> torch.Tensor:
+    """
+    Convert a skew-symmetric matrix to a rotation vector in a differentiable way
+    [
+        [ 0, -z,  y],
+        [ z,  0, -x],
+        [-y,  x,  0],
+    ]
+    Args:
+        skew_symmetric: Skew-symmetric matrix of shape [..., 3, 3]
+    Returns:
+        torch.Tensor of shape [..., 3]
+    """
+    assert is_rotmat(skew_symmetric), skew_symmetric.shape
+    rotvec = torch.stack(
+        (
+            skew_symmetric[..., 2, 1] - skew_symmetric[..., 1, 2],
+            skew_symmetric[..., 0, 2] - skew_symmetric[..., 2, 0],
+            skew_symmetric[..., 1, 0] - skew_symmetric[..., 0, 1],
+        ),
+        dim=-1,
+    )
+    rotvec = rotvec / 2.0
+    return rotvec
+
+
+def integrate_rotmat(
+    rt: torch.Tensor, dr_dt: torch.Tensor, dt: float | torch.Tensor, body_frame: bool = True
+) -> torch.Tensor:
+    """
+    Integrate a rotation matrix `rt` by the derivative `dr_dt` over the time interval `dt`.
+    Args:
+        rt: Rotation matrix, shape [..., 3, 3]
+        dr_dt: Derivative of the rotation matrix, shape [..., 3, 3]
+        dt: Time interval to integrate over, scalar or a tensor of shape () or [..., 1]
+        body_frame: If True, the integration is done in the body frame (post-multiply),
+            otherwise in the inertial frame (pre-multiply).
+    Returns:
+        Integrated unit quaternion, shape [..., 4]
+    """
+    assert rt.shape == dr_dt.shape, f'{rt.shape} != {dr_dt.shape}'
+    assert is_rotmat(rt), f'{rt.shape} not a rotation matrix'
+    is_3x3 = is_rotmat_3x3(rt)
+    if not is_3x3:
+        rt = rotmat_as_3x3(rt)
+        dr_dt = rotmat_as_3x3(dr_dt)
+    if isinstance(dt, torch.Tensor):
+        assert dt.ndim in (0, rt.ndim, rt.ndim - 1), f'dt.ndim = {dt.ndim} | {rt.ndim} | {rt.ndim - 1}'
+        if dt.ndim == rt.ndim:
+            assert dt.shape[-2:] == (1, 1), dt.shape
+            dt = dt.squeeze(-1)
+    if body_frame:
+        omega = skew_symmetric_to_rotvec(rotmat_inverse(rt) @ dr_dt)
+    else:
+        omega = skew_symmetric_to_rotvec(dr_dt @ rotmat_inverse(rt))
+    dr = roma.rotvec_to_rotmat(omega * dt)
+    if body_frame:
+        rt = rt @ dr
+    else:
+        rt = dr @ rt
+    if not is_3x3:
+        rt = rotmat_as_9(rt)
+    return rt
+
+
+def integrate_rotation(
+    rt: torch.Tensor,
+    dr_dt: torch.Tensor,
+    dt: float | torch.Tensor,
+    body_frame: bool = True,
+    half_cover: bool = True,
+) -> torch.Tensor:
+    """
+    Integrate the rotation `rt` by the derivative `dr_dt` over the time interval `dt` on the SO(3) manifold.
+    """
+    if is_quaternion(rt):
+        return integrate_unitquat(rt, dr_dt, dt, body_frame=body_frame, half_cover=half_cover)
+    if is_rotmat(rt):
+        return integrate_rotmat(rt, dr_dt, dt, body_frame=body_frame)
+    raise NotImplementedError(f'integrate_rotation not yet implemented for format {rt.shape}')
+
+
+class PiZeroFlowMatchingModule(ConfigurableModule[PiZeroFlowMatchingModuleConfig]):
+    def __init__(self, config: PiZeroFlowMatchingModuleConfig, control_tokenizer: EmptyTokenizer):
+        super().__init__(config)
+        del control_tokenizer
+        self.noised_control_proj = NoisedControlProjector(self.config.noised_control_proj_config)
+        self.robot_state_proj = RobotStateProjector(self.config.robot_state_proj_config)
+        self.control_decoder = PiZeroFlowMatchingDecoder(config=self.config.control_decoder_config)
+        self.output_proj = make_mlp(
+            [self.config.token_size, 3 + self.config.rotation_components + 1],
+            activation=torch.nn.GELU,
+            activate_final=False,
+        )
+
+    def forward(
+        self, vlm_input: RoboticsFlowInput, vlm_output: VLMOutput, cache: Optional[transformers.Cache] = None
+    ) -> RoboticsOutput:
+        robot_state_tokens = self.robot_state_proj(vlm_input)
+        noised_tokens = self.noised_control_proj(vlm_input.flow_input)
+        output_tokens = self.control_decoder(
+            control_tokens=noised_tokens,
+            robot_state_tokens=robot_state_tokens,
+            llm_kv_tokens=vlm_output.llm_output.past_key_values,
+            attn_mask=vlm_input.attn_mask,
+            cache=cache,
+        )
+        contols = self.output_proj(output_tokens)
+        (translation, rotation, gripper) = torch.split(
+            contols, [3, self.config.rotation_components, 1], dim=-1
+        )
+        return RoboticsOutput.make_empty().replace(
+            translation=translation, rotation=rotation, gripper=gripper
+        )
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        vlm_input: RoboticsFlowInput,
+        vlm_output: VLMOutput,
+        processor: PiZeroFlowMatchingProcessor,
+        use_cache: bool = True,
+        **kwargs,
+    ) -> RoboticsOutput:
+        del kwargs
+        (batch_size, vlm_seq_len) = vlm_input.input_ids.shape[:2]
+        device = vlm_input.input_ids.device
+        if use_cache:
+            max_cache_len = (
+                vlm_seq_len
+                + processor.config.control_io_config.future_controls_sequence_length
+                + processor.config.control_io_config.past_scalars_sequence_length
+            )
+            cache = transformers.StaticCache(
+                config=transformers.PretrainedConfig(
+                    head_dim=self.config.control_decoder_config.block_config.head_dim,
+                    num_key_value_heads=self.config.control_decoder_config.block_config.num_kv_heads,
+                    num_hidden_layers=self.config.control_decoder_config.num_blocks,
+                ),
+                max_batch_size=batch_size,
+                max_cache_len=max_cache_len,
+                device=device,
+            )
+        else:
+            cache = None
+        flow_input: FlowInput = processor.sample_t0_input(batch_size=batch_size, device=device)
+        step_size = 1 / processor.config.num_inference_steps
+        translation = flow_input.translation_t0
+        rotation = flow_input.rotation_t0
+        gripper = flow_input.gripper_t0
+        vlm_input = vlm_input.replace(
+            **{
+                'flow_input.timestep': flow_input.timestep,
+                'flow_input.translation_t': translation,
+                'flow_input.rotation_t': rotation,
+                'flow_input.gripper_t': gripper,
+            }
+        )
+        for _ in range(processor.config.num_inference_steps):
+            model_output: RoboticsOutput = self(vlm_input, vlm_output, cache)
+            translation = translation + step_size * model_output.translation
+            rotation = integrate_rotation(rt=rotation, dr_dt=model_output.rotation, dt=step_size)
+            gripper = gripper + step_size * model_output.gripper
+            timestep = vlm_input.flow_input.timestep + step_size
+            if processor.config.rotation_format == RotationFormat.QUATERNION:
+                rotation = quaternion_half_cover(rotation)
+            vlm_input = vlm_input.replace(
+                **{
+                    'flow_input.timestep': timestep,
+                    'flow_input.translation_t': translation,
+                    'flow_input.rotation_t': rotation,
+                    'flow_input.gripper_t': gripper,
+                }
+            )
+        output = RoboticsOutput.make_empty().replace(
+            translation=translation, rotation=rotation, gripper=gripper
+        )
+        return output
+
+    @property
+    def fsdp_wrap_modules(self) -> Dict[torch.nn.Module, Dict[str, Any]]:
+        return self.control_decoder.fsdp_wrap_modules | {
+            self: {},
+            self.robot_state_proj: {},
+            self.noised_control_proj: {},
+            self.output_proj: {},
+        }
+
+
+class VLAM(ConfigurableModule[VLAMConfig]):
+    def __init__(self, config: VLAMConfig):
+        super().__init__(config)
+        self.vlm = Qwen3VL(config=self.config.vlm_config)
+        self.processor = PiZeroFlowMatchingProcessor(
+            config=self.config.processor_config, vlm_processor=self.vlm.processor
+        )
+        self.control_module = PiZeroFlowMatchingModule(
+            config=self.config.control_module_config, control_tokenizer=self.processor.control_tokenizer
+        )
+
+    def forward(
+        self,
+        inputs: RoboticsInput,
+        use_cache: Optional[bool] = True,
+        output_hidden_states: Optional[bool] = None,
+    ) -> RoboticsOutput:
+        del output_hidden_states
+        vlm_output = self.vlm(inputs=inputs, use_cache=use_cache, output_hidden_states=True)
+        control_output = self.control_module(vlm_input=inputs, vlm_output=vlm_output)
+        output = control_output.replace(llm_output=vlm_output.llm_output)
+        return output
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        inputs: RoboticsInput,
+        use_cache: Optional[bool] = True,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+    ) -> RoboticsOutput:
+        del output_hidden_states
+        vlm_output = self.vlm(
+            inputs=inputs, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=True
+        )
+        control_output = self.control_module.generate(
+            vlm_input=inputs, vlm_output=vlm_output, processor=self.processor
+        )
+        output = control_output.replace(llm_output=vlm_output.llm_output)
+        return output
+
+    @property
+    def fsdp_wrap_modules(self) -> Dict[torch.nn.Module, Dict[str, Any]]:
+        return {
+            **self.vlm.fsdp_wrap_modules,
+            **self.control_module.fsdp_wrap_modules,
+            self.vlm: {},
+            self.control_module: {},
+        }
+
+
+MainModel = VLAM

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/hf_export/keen-fuchsia-mandrill/src/processing_pizero_fm_qwen3_vl.py

@@ -0,0 +1,1955 @@
+import warnings
+from abc import abstractmethod
+from functools import cached_property
+from typing import Dict, List, Optional, Tuple, TypeVar
+
+import numpy as np
+import PIL.Image
+import roma
+import torch
+import torchvision.transforms.v2
+import transformers
+from databib.config import Configurable
+from databib.template import Template
+
+from .common_pizero_fm_qwen3_vl import (
+    FlowInput,
+    ReferenceFrame,
+    ResizeMode,
+    RoboticsControlPlan,
+    RoboticsFlowInput,
+    RoboticsInput,
+    RoboticsOutput,
+    RoboticsTarget,
+    RotationFormat,
+    expand_dims,
+    is_quaternion,
+    is_rotmat,
+    is_rotmat_3x3,
+    is_rotmat_9,
+    quaternion_half_cover,
+    rotmat_as_3x3,
+    rotmat_as_9,
+    rotmat_inverse,
+)
+from .configuration_pizero_fm_qwen3_vl import (
+    BoundsNormalizerConfig,
+    ControlDataIOConfig,
+    ControlTokenizerConfig,
+    DatasetStatsNormalizerConfig,
+    EmptyTokenizerConfig,
+    IdentityNormalizerConfig,
+    ImageSizeConfig,
+    NormalizerConfig,
+    PiZeroFlowProcessorConfig,
+    RegressionProcessorConfig,
+    RotationStereomapNormalizerConfig,
+    VLAMProcessorConfig,
+    VLMProcessorConfig,
+)
+
+ControlTokenizerConfigT = TypeVar('ControlTokenizerConfigT', bound=ControlTokenizerConfig)
+
+
+class ControlTokenizer(Configurable[ControlTokenizerConfigT], Template[ControlTokenizerConfigT]):
+    @abstractmethod
+    def __call__(self, *args, **kwargs) -> str:
+        """Given GT actions and possibly other information, output text control. Gets appened to the prompt"""
+
+
+class EmptyTokenizer(ControlTokenizer[EmptyTokenizerConfig]):
+    """
+    Takes the LLM hidden states from `llm_layer_indices` and concatenates them to produce the
+    desired result. Includes the hidden states for the image tokens.
+    """
+
+    def __init__(self, config, tokenizer: transformers.PreTrainedTokenizerBase) -> None:
+        super().__init__(config)
+        self.tokenizer = tokenizer
+
+    def __call__(self, *_) -> str:
+        return ''
+
+
+NormalizerConfigT = TypeVar('NormalizerConfigT', bound=NormalizerConfig)
+
+
+class Normalizer(Configurable[NormalizerConfigT], Template[NormalizerConfigT]):
+    @abstractmethod
+    def normalize(self, value: torch.Tensor, **kwargs) -> torch.Tensor:
+        """
+        Normalize the input value.
+
+        Args:
+            value: Tensor to be normalized
+            **kwargs: Implmentation-specific arguments for normalization
+        Returns:
+            Normalized tensor of the same shape as input
+        """
+
+    @abstractmethod
+    def unnormalize(self, value: torch.Tensor, **kwargs) -> torch.Tensor:
+        """
+        Unnormalize the input value.
+
+        Args:
+            value: Tensor to be normalized
+            **kwargs: Implmentation-specific arguments for normalization
+        Returns:
+            Unnormalized tensor of the same shape as input
+        """
+
+
+class IdentityNormalizer(Normalizer[IdentityNormalizerConfig]):
+    def normalize(self, value: torch.Tensor, **kwargs) -> torch.Tensor:
+        del kwargs
+        return value
+
+    def unnormalize(self, value: torch.Tensor, **kwargs) -> torch.Tensor:
+        del kwargs
+        return value
+
+
+def np_unique(data: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Compute unique elements in data and corresponding indices.
+
+    np.unique returns the values in a sorted order, even if the source is not sorted. Thus, if you simply
+    run np.unique on unsorted data, the indices you will get will be invalid.
+
+    """
+    (_, indices, inverse) = np.unique(data, return_index=True, return_inverse=True)
+    (_, indices_of_first_occurence, inverse_indices, counts) = np.unique(
+        indices[inverse], return_index=True, return_inverse=True, return_counts=True
+    )
+    unique_ids = data[indices_of_first_occurence]
+    return unique_ids, indices_of_first_occurence, inverse_indices, counts
+
+
+def _broadcast_shapes(
+    value: torch.Tensor, low: torch.Tensor, high: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Broadcast shapes for normalization:
+    Args:
+        value: torch.Tensor of shape [..., num_components]. The entire shape might be:
+            - [num_components]: `value` has no batch dimension
+            - [num_datasets, num_components]: `value` contains entries *aligned* with the dataset bounds
+                contained in `low` and `high`
+            - [num_datasets, ..., num_components]: `value` contains entries *aligned* with the dataset bounds
+                contained in `low` and `high`
+            - [..., num_components]: `value` contains multiple dimensions. In this case, `low` and `high`
+                must be for a single dataset, i.e. `num_datasets = 1`
+
+        low: torch.Tensor, shape [num_datasets, num_components], where `num_datasets` can be 1 when `low`
+            contains normalization bounds for a single dataset
+        high: torch.Tensor, shape [num_datasets, num_components], where `num_datasets` can be 1 when `high`
+            contains normalization bounds for a single dataset
+    Returns:
+        Tuple of torch.Tensors (low, high), where `low` and `high` have the same number of dimensions as `value`
+    """
+    assert low.ndim == high.ndim == 2, f'{low.shape} != {high.shape} or ndim != 2'
+    assert value.shape[-1] == low.shape[-1] == high.shape[-1], f'{value.shape} != {low.shape} / {high.shape}'
+    if value.ndim == low.ndim == high.ndim:
+        return low, high
+    if value.ndim < low.ndim:
+        assert low.ndim == high.ndim == 2, f'{low.shape}, {high.shape}'
+        assert low.shape[0] == high.shape[0] == 1, f'{low.shape}, {high.shape}'
+        (low, high) = (low.view(-1), high.view(-1))
+        return low, high
+    if low.shape[0] == high.shape[0] == 1:
+        low = expand_dims(low.view(-1), ndim=value.ndim, order=[-1, 1])
+        high = expand_dims(high.view(-1), ndim=value.ndim, order=[-1, 1])
+    else:
+        assert value.shape[0] == low.shape[0] == high.shape[0], f'{value.shape} != {low.shape} / {high.shape}'
+        low = expand_dims(low, ndim=value.ndim, order=[1, -1, 1])
+        high = expand_dims(high, ndim=value.ndim, order=[1, -1, 1])
+    return low, high
+
+
+def normalize_gripper_by_bounds(
+    value: torch.Tensor, low: torch.Tensor, high: torch.Tensor, binary: bool = True
+) -> torch.Tensor:
+    """
+    If binary, normalize to [0, 1], otherwise normalize to [-1, 1]
+    """
+    (low, high) = _broadcast_shapes(value, low, high)
+    (low, high) = (low.to(device=value.device), high.to(device=value.device))
+    if binary:
+        return torch.clamp((value - low) / torch.clamp(high - low, min=1e-08), min=0.0, max=1.0)
+    return torch.clamp(2 * (value - low) / torch.clamp(high - low, min=1e-08) - 1, min=-1.0, max=1.0)
+
+
+def unnormalize_by_moments(value: torch.Tensor, mean: torch.Tensor, std: torch.Tensor) -> torch.Tensor:
+    (mean, std) = _broadcast_shapes(value, mean, std)
+    (mean, std) = (mean.to(device=value.device), std.to(device=value.device))
+    return value * (std + 1e-08) + mean
+
+
+def normalize_by_moments(value: torch.Tensor, mean: torch.Tensor, std: torch.Tensor) -> torch.Tensor:
+    (mean, std) = _broadcast_shapes(value, mean, std)
+    (mean, std) = (mean.to(device=value.device), std.to(device=value.device))
+    return (value - mean) / (std + 1e-08)
+
+
+def unnormalize_by_bounds(value: torch.Tensor, low: torch.Tensor, high: torch.Tensor) -> torch.Tensor:
+    (low, high) = _broadcast_shapes(value, low, high)
+    (low, high) = (low.to(device=value.device), high.to(device=value.device))
+    return 0.5 * (value + 1) * (high - low) + low
+
+
+def normalize_by_bounds(value: torch.Tensor, low: torch.Tensor, high: torch.Tensor) -> torch.Tensor:
+    (low, high) = _broadcast_shapes(value, low, high)
+    (low, high) = (low.to(device=value.device), high.to(device=value.device))
+    return torch.clamp(2 * (value - low) / torch.clamp(high - low, min=1e-08) - 1, min=-1.0, max=1.0)
+
+
+class DatasetStatsNormalizer(Normalizer[DatasetStatsNormalizerConfig]):
+    def __init__(self, config: DatasetStatsNormalizerConfig):
+        super().__init__(config)
+        self._norm_stats = self._load_norm_stats()
+
+    def _load_norm_stats(self) -> Dict[str, Dict[str, Dict[str, torch.Tensor]]]:
+        norm_stats = {
+            'austin_buds_dataset': {
+                'low': [0.3499317765235901, -0.2854413390159607, 0.010516085661947727],
+                'high': [0.7243335843086243, 0.20652863383293152, 0.3218296766281128],
+            },
+            'austin_sailor_dataset': {
+                'low': [0.387094110250473, -0.3164229393005371, 0.024492919445037842],
+                'high': [0.6869593262672424, 0.2086469978094101, 0.2551962733268738],
+            },
+            'austin_sirius_dataset': {
+                'low': [0.0, -0.11814527958631516, 0.0],
+                'high': [0.532875120639801, 0.26084619760513306, 0.27225059270858765],
+            },
+            'bc_z': {
+                'low': [-0.3956047296524048, -0.11924505233764648, 0.601338267326355],
+                'high': [0.332028865814209, 0.3088575601577759, 0.98329097032547],
+            },
+            'berkeley_autolab_ur5': {
+                'low': [0.3020566999912262, -0.21297279000282288, -0.18836002051830292],
+                'high': [0.6132073998451233, 0.30656182765960693, 0.12212439626455307],
+            },
+            'berkeley_cable_routing': {
+                'low': [0.4641263782978058, -0.2806571424007416, 0.030183622613549232],
+                'high': [0.6452807784080505, 0.28204888105392456, 0.1557157188653946],
+            },
+            'berkeley_fanuc_manipulation': {
+                'low': [0.3718133866786957, -0.4071895182132721, 0.01847645826637745],
+                'high': [0.7200658321380615, 0.3128541111946106, 0.5413243770599365],
+            },
+            'bridge': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_32b': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_coarse_max3': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_full_tread_8b_k5': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_hindsight': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_orig': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_nils': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_paraphrase_k10': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_paraphrase_k5': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_paraphrase_k5_mix50': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_rich_properties': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_rich_properties_full': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_rich_properties_mix50': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_rich_properties_p30': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_rich_properties_p50': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_steering': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_tread': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_tread_full': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'bridge_tread_k10': {
+                'low': [0.1711955964565277, -0.15639324486255646, -0.048255354166030884],
+                'high': [0.4604376256465912, 0.24112474918365479, 0.18886254727840424],
+            },
+            'cmu_stretch': {
+                'low': [0.017430847510695457, 0.0, 0.46050605177879333],
+                'high': [0.33094948530197144, 0.0, 1.0952961444854736],
+            },
+            'dlr_edan_shared_control': {
+                'low': [-0.729511022567749, 0.077408567070961, 0.2658006250858307],
+                'high': [-0.13719859719276428, 0.5719971060752869, 0.7898909449577332],
+            },
+            'droid': {
+                'low': [0.26669958233833313, -0.43774399161338806, -0.048167888075113297],
+                'high': [0.7774086594581604, 0.42832574248313904, 0.7760910391807556],
+            },
+            'fmb': {
+                'low': [0.3655048608779907, -0.28729698061943054, 0.033201027661561966],
+                'high': [0.6782684326171875, 0.209969624876976, 0.3331448435783386],
+            },
+            'fractal20220817_data': {
+                'low': [0.3249714970588684, -0.2818704843521118, 0.1410011649131775],
+                'high': [0.8754204511642456, 0.21279653906822205, 1.071526288986206],
+            },
+            'furniture_bench_dataset': {
+                'low': [0.36915361881256104, -0.180975541472435, 0.0058300793170928955],
+                'high': [0.6652880311012268, 0.1772783100605011, 0.18316447734832764],
+            },
+            'iamlab_cmu_pickup_insert': {
+                'low': [0.31449857354164124, -0.20315787196159363, 0.06785127520561218],
+                'high': [0.6472027897834778, 0.20840713381767273, 0.3700340986251831],
+            },
+            'jaco_play': {
+                'low': [-0.3789186179637909, -0.6194459795951843, 0.16865813732147217],
+                'high': [0.21203258633613586, -0.26914602518081665, 0.38958534598350525],
+            },
+            'kuka': {
+                'low': [0.4765772819519043, -0.14815208315849304, 0.06674224138259888],
+                'high': [0.6515637040138245, 0.2447487711906433, 0.28018367290496826],
+            },
+            'language_table': {
+                'low': [0.19237099587917328, -0.2962527573108673, 0.0],
+                'high': [0.6171894669532776, 0.30645298957824707, 0.0],
+            },
+            'nyu_franka_play_dataset': {
+                'low': [0.13936959207057953, 0.07645522058010101, 0.19364508986473083],
+                'high': [0.5920727252960205, 0.6584802269935608, 0.8056891560554504],
+            },
+            'roboset': {
+                'low': [0.18437016010284424, -0.25699371099472046, 0.15134164690971375],
+                'high': [0.543661892414093, 0.29646238684654236, 0.6682320833206177],
+            },
+            'roboturk': {
+                'low': [0.28454264998435974, -0.3288349509239197, -0.09349551796913147],
+                'high': [0.8773894309997559, 0.2857522964477539, 0.32863926887512207],
+            },
+            'stanford_hydra_dataset': {
+                'low': [0.23737286031246185, -0.26521679759025574, 0.09069013595581055],
+                'high': [0.7124238014221191, 0.25299057364463806, 0.49505406618118286],
+            },
+            'taco_play': {
+                'low': [0.1368357390165329, -0.4297449290752411, 0.20516259968280792],
+                'high': [0.6700438857078552, 0.5943909883499146, 0.5966404676437378],
+            },
+            'toto': {
+                'low': [-0.09177927672863007, -0.3571659028530121, 0.2196546494960785],
+                'high': [0.6757593750953674, 0.2889021635055542, 0.5011094212532043],
+            },
+            'ucsd_kitchen_dataset': {
+                'low': [0.18739914894104004, -0.18234309554100037, 0.04897069185972214],
+                'high': [0.6410437822341919, 0.20632223784923553, 0.5983893275260925],
+            },
+            'utaustin_mutex': {
+                'low': [0.3217194080352783, -0.4733337163925171, 0.014122226275503635],
+                'high': [0.5321439504623413, 0.3733823001384735, 0.5785381197929382],
+            },
+            'viola': {
+                'low': [0.40061360597610474, -0.25196850299835205, 0.010269512422382832],
+                'high': [0.6458418369293213, 0.17776551842689514, 0.4456312954425812],
+            },
+        }
+        return {
+            dataset_name: {
+                key: torch.tensor(value, dtype=torch.float32) for (key, value) in dataset_stats.items()
+            }
+            for (dataset_name, dataset_stats) in norm_stats.items()
+        }
+
+    def _broadcast_norm_stats_to_dataset_name(
+        self, dataset_name: np.ndarray
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Create an array of normalization bounds corresponding to dataset names
+        Args:
+            dataset_name: Array of shape [B] of dataset names for which to fetch normalization stats.
+                Note the values can be repeated
+        Returns:
+            Tuple of (low, high) or (norm, std) stats, each of shape [B, -1]
+        """
+        if self.config.mode == 'mean':
+            (stats_key_1, stats_key_2) = ('mean', 'std')
+        else:
+            (stats_key_1, stats_key_2) = ('low', 'high')
+        (unique_names, _, inverse_indices, _) = np_unique(dataset_name)
+        stats_1 = np.zeros([len(unique_names), self._component_size], dtype=np.float32)
+        stats_2 = np.zeros([len(unique_names), self._component_size], dtype=np.float32)
+        for i, ds_name in enumerate(unique_names):
+            stats_1[i] = self._norm_stats[ds_name][stats_key_1].numpy()
+            stats_2[i] = self._norm_stats[ds_name][stats_key_2].numpy()
+        stats_1 = stats_1[inverse_indices]
+        stats_2 = stats_2[inverse_indices]
+        return torch.from_numpy(stats_1), torch.from_numpy(stats_2)
+
+    @property
+    def _component_size(self) -> int:
+        return list(list(self._norm_stats.values())[0].values())[0].shape[-1]
+
+    def normalize(self, value: torch.Tensor, dataset_name: np.ndarray, **kwargs) -> torch.Tensor:
+        del kwargs
+        if self.config.mode == 'mean':
+            (mean, std) = self._broadcast_norm_stats_to_dataset_name(dataset_name)
+            output = normalize_by_moments(value, mean=mean, std=std)
+        else:
+            (low, high) = self._broadcast_norm_stats_to_dataset_name(dataset_name)
+            output = normalize_by_bounds(value, low=low, high=high)
+        return output
+
+    def unnormalize(self, value: torch.Tensor, dataset_name: np.ndarray, **kwargs) -> torch.Tensor:
+        del kwargs
+        if self.config.mode == 'mean':
+            (mean, std) = self._broadcast_norm_stats_to_dataset_name(dataset_name)
+            output = unnormalize_by_moments(value, mean=mean, std=std)
+        else:
+            (low, high) = self._broadcast_norm_stats_to_dataset_name(dataset_name)
+            output = unnormalize_by_bounds(value, low=low, high=high)
+        return output
+
+
+class BoundsNormalizer(Normalizer[BoundsNormalizerConfig]):
+    def __init__(self, config: BoundsNormalizerConfig):
+        super().__init__(config)
+        self.low = torch.tensor(self.config.low, dtype=torch.float32).view(1, -1)
+        self.high = torch.tensor(self.config.high, dtype=torch.float32).view(1, -1)
+
+    def normalize(self, value: torch.Tensor, **kwargs) -> torch.Tensor:
+        del kwargs
+        return normalize_by_bounds(value, low=self.low, high=self.high)
+
+    def unnormalize(self, value: torch.Tensor, **kwargs) -> torch.Tensor:
+        del kwargs
+        return unnormalize_by_bounds(value, low=self.low, high=self.high)
+
+
+def euler_to_rotmat(angles: torch.Tensor) -> torch.Tensor:
+    """
+    Args:
+        angles: Euler angles in radians in the format 'xyz', shape [..., 3]
+    Returns:
+        torch.Tensor of shape [..., 3, 3] containing rotation matrices
+    """
+    return roma.euler_to_rotmat(convention='xyz', angles=angles, degrees=False)
+
+
+def euler_to_unit_quaternion(angles: torch.Tensor) -> torch.Tensor:
+    """
+    Args:
+        angles: Euler angles in radians in the format 'xyz', shape [..., 3]
+    Returns:
+        torch.Tensor of shape [..., 4] containing unit quaternions
+    """
+    return roma.euler_to_unitquat(convention='xyz', angles=angles, degrees=False, normalize=True)
+
+
+def normalize_quaternion(quaternion: torch.Tensor, eps: float = 1e-08) -> torch.Tensor:
+    """
+    Args:
+        quaternion: Unnormalized quaternion, torch.Tensor of shape [..., 4]
+        eps: Small constant to prevent division by zero
+    Returns:
+        torch.Tensor of shape [..., 4] of unit quaternions
+    """
+    return quaternion / (quaternion.norm(dim=-1, keepdim=True).detach() + eps)
+
+
+def quaternion_to_euler(quaternion: torch.Tensor) -> torch.Tensor:
+    """
+    Args:
+        quaternion: torch.Tensor of shape [..., 4]; Can be non-normalized
+    Returns:
+        torch.Tensor of shape [..., 3, 3] containing rotation matrices in SO(3)
+    """
+    unit_quat = normalize_quaternion(quaternion)
+    rotmat = roma.unitquat_to_euler(convention='xyz', quat=unit_quat, as_tuple=False, degrees=False)
+    return rotmat
+
+
+def quaternion_to_rotmat(quaternion: torch.Tensor) -> torch.Tensor:
+    """
+    Args:
+        quaternion: torch.Tensor of shape [..., 4]; Can be non-normalized
+    Returns:
+        torch.Tensor of shape [..., 3, 3] containing rotation matrices in SO(3)
+    """
+    unit_quat = normalize_quaternion(quaternion)
+    rotmat = roma.unitquat_to_rotmat(unit_quat)
+    return rotmat
+
+
+def rotmat_to_unit_quaternion(rotmat: torch.Tensor) -> torch.Tensor:
+    """
+    Args:
+        rotmat: Batch of rotation matrices, shape [..., 3, 3]
+    Returns:
+        Batch of unit quaternions, shape [..., 4]
+    """
+    rotmat = rotmat_as_3x3(rotmat)
+    return roma.rotmat_to_unitquat(rotmat)
+
+
+def rotmat_to_euler(rotmat: torch.Tensor) -> torch.Tensor:
+    """
+    Args:
+        rotmat: Batch of rotation matrices, shape [..., 3, 3]
+    Returns:
+        Batch of Euler angles in radiant, shape [..., 3]
+    """
+    rotmat = rotmat_as_3x3(rotmat)
+    return roma.rotmat_to_euler(convention='xyz', rotmat=rotmat, as_tuple=False, degrees=False)
+
+
+def symmetric_orthogonalization(x: torch.Tensor) -> torch.Tensor:
+    """
+    Maps 9D input vectors onto SO(3) via symmetric orthogonalization.
+        - Let SVD(M) = U \Sigma V^T
+        - Returned value is SVD+(M) =  U diag(1, 1, det(UV^T)) V^T
+        - det(UV^T) ensures that det(SVD+(M)) = 1
+        - The return value is a rotation matrix (ortonormal) with the least-squares distance to M
+
+    Args:
+        x: Input matrices, not necessarily orthonormal, shape [..., 9] or [..., 3, 3]
+    Returns:
+        torch.Tensor with the same shape as x, where each inner 3x3 matrix is in SO(3)
+    """
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            'ignore', message='In CPU autocast, but the target dtype is not supported. Disabling autocast.'
+        )
+        with torch.autocast(device_type=x.device.type, dtype=torch.float32):
+            matrices = x.view(-1, 3, 3)
+            matrices = matrices.to(dtype=torch.float32)
+            (u, s, v) = torch.svd(matrices)
+            vt = torch.transpose(v, 1, 2)
+            det = torch.det(torch.matmul(u, vt)).view(-1, 1, 1)
+            diag_vt = torch.cat((vt[:, :2, :], vt[:, -1:, :] * det), dim=1)
+            result = torch.matmul(u, diag_vt)
+            result = result.view(*x.shape)
+    result = result.to(dtype=x.dtype)
+    return result
+
+
+def is_rotmat_orthonormal(
+    rotmat: torch.Tensor, epsilon: float = 1e-06, reduction: str = 'none'
+) -> torch.Tensor | bool:
+    """
+    Check if a rotation matrix is orthonormal or not.
+    Args:
+        rotmat: torch.Tensor of shape [..., 3, 3] or [..., 9]
+        epsilon: Tolerance for numerical comparisons. Bigger values allow for more freedom. Generally,
+            anything smaller than 1e-6 might incorrectly detect some otrhonormal matrices as not
+        reduction:
+            'none' - returns torch.Tensor of bools with the same batch shape
+            'all' - returns a bool, True is ALL matrices in the batch are orthonormal
+    Returns:
+        torch.Tensor with the same batch shape or bool
+    """
+    assert is_rotmat(rotmat)
+    rotmat = rotmat_as_3x3(rotmat.to(dtype=torch.float32))
+    is_orthonormal = roma.is_orthonormal_matrix(rotmat, epsilon=epsilon)
+    if reduction == 'none':
+        return is_orthonormal
+    if reduction == 'all':
+        return bool(torch.all(is_orthonormal).item())
+    raise ValueError(f'Unknown reduction mode {reduction}')
+
+
+def is_orthonormal_rotmat(rotmat: torch.Tensor, epsilon=0.01, reduction='none') -> bool:
+    """
+    Checks if the tensor shape matches that of a rotmat. If the last dimensions of shape are 3x3,
+    also checks if the data is a valid rotmat. This is to avoid a possible clash with euler angles
+    when accidentally `rotmat.shape[-2:] == [3, 3]`
+    """
+    return (
+        is_rotmat_9(rotmat)
+        or is_rotmat_3x3(rotmat)
+        and is_rotmat_orthonormal(rotmat, epsilon=epsilon, reduction=reduction)
+    )
+
+
+def is_euler(euler: torch.Tensor) -> bool:
+    return euler.shape[-1] == 3 and not is_orthonormal_rotmat(euler, reduction='all')
+
+
+def normalize_rotation(rotation: torch.Tensor) -> torch.Tensor:
+    if is_quaternion(rotation):
+        return normalize_quaternion(rotation)
+    if is_euler(rotation):
+        return rotation
+    if is_rotmat(rotation):
+        is_flat = is_rotmat_9(rotation)
+        rotation = rotmat_as_3x3(rotation) if is_flat else rotation
+        rotmat = roma.special_gramschmidt(rotation)
+        rotmat = rotmat_as_9(rotmat) if is_flat else rotmat
+        return rotmat
+    raise ValueError(f'Unknown rotation format: {rotation.shape}')
+
+
+def rotation_format_from_tensor(rotation) -> RotationFormat:
+    if is_quaternion(rotation):
+        return RotationFormat.QUATERNION
+    if is_orthonormal_rotmat(rotation, reduction='all'):
+        return RotationFormat.ROTMAT
+    if is_euler(rotation):
+        return RotationFormat.EULER
+    raise ValueError(f'Tensor shape {rotation.shape} is not a valid rotation format')
+
+
+def is_unit_quaternion(
+    quaternion: torch.Tensor, epsilon: float = 1e-08, reduction: str = 'none'
+) -> torch.Tensor | bool:
+    """
+    Check if a quternion is normalized or not.
+    Args:
+        quaternion: torch.Tensor of shape [..., 4]
+        tolerance: Tolerance for numerical comparisons
+        reduction:
+            'none' - returns torch.Tensor of bools with the same batch shape
+            'all' - returns a bool, True if ALL quaternions in the batch are normalized
+    Returns:
+        torch.Tensor with the same batch shape or bool
+    """
+    if not is_quaternion(quaternion):
+        return False
+    is_norm = torch.isclose(
+        quaternion.norm(dim=-1, keepdim=True),
+        torch.tensor(1.0, dtype=quaternion.dtype, device=quaternion.device),
+        atol=epsilon,
+    )
+    if reduction == 'none':
+        return is_norm
+    if reduction == 'all':
+        return bool(torch.all(is_norm).item())
+    raise ValueError(f'Unknown reduction mode {reduction}')
+
+
+def convert_rotation(
+    rotation: torch.Tensor | np.ndarray,
+    output_format: RotationFormat,
+    autonorm: bool = True,
+    half_cover: bool = True,
+) -> torch.Tensor | np.ndarray:
+    is_np = isinstance(rotation, np.ndarray)
+    if is_np:
+        rotation = torch.from_numpy(rotation)
+    if is_quaternion(rotation):
+        if autonorm and not is_unit_quaternion(rotation, reduction='all'):
+            rotation = normalize_quaternion(rotation)
+        if output_format == RotationFormat.QUATERNION:
+            output = rotation
+        elif output_format == RotationFormat.ROTMAT:
+            output = rotmat_as_9(quaternion_to_rotmat(rotation))
+        elif output_format == RotationFormat.EULER:
+            output = quaternion_to_euler(rotation)
+        else:
+            raise NotImplementedError(f'Unsupported rotation format: {output_format}')
+    elif is_orthonormal_rotmat(rotation, reduction='all'):
+        if autonorm and not is_rotmat_orthonormal(rotation, epsilon=0.01, reduction='all'):
+            rotation = symmetric_orthogonalization(rotation)
+        if output_format == RotationFormat.QUATERNION:
+            output = rotmat_to_unit_quaternion(rotation)
+        elif output_format == RotationFormat.ROTMAT:
+            output = rotmat_as_9(rotation)
+        elif output_format == RotationFormat.EULER:
+            output = rotmat_to_euler(rotation)
+        else:
+            raise NotImplementedError(f'Unsupported rotation format: {output_format}')
+    elif is_euler(rotation):
+        if output_format == RotationFormat.QUATERNION:
+            output = euler_to_unit_quaternion(rotation)
+        elif output_format == RotationFormat.ROTMAT:
+            output = rotmat_as_9(euler_to_rotmat(rotation))
+        elif output_format == RotationFormat.EULER:
+            output = rotation
+        else:
+            raise NotImplementedError(f'Unsupported rotation format: {output_format}')
+    else:
+        raise ValueError(f'Unknown rotation encoding with shape {rotation.shape}')
+    if output_format == RotationFormat.QUATERNION and half_cover:
+        output = quaternion_half_cover(output)
+    if is_np:
+        output = output.numpy()
+    return output
+
+
+def apply_rotation(rotation: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+    """
+    Rotate `value` by `rotation`
+    Args:
+        rotation: torch.Tensor, euler, quaternion or rotmat. Any batch shape that can be expanded
+            such that it broadcasts to `value`
+        value: torch.Tensor. Supported shapes:
+            - Rotmat: [B, ..., 3, 3] or [B, ..., 9]
+            - Quaternion: [B, ..., 4]
+            - 3D vector: [B, ..., 3]
+    Returns:
+        torch.Tensor of the same shape as `value`
+    """
+    rotation = rotmat_as_3x3(convert_rotation(rotation, RotationFormat.ROTMAT))
+    quaternion = is_quaternion(value)
+    if quaternion:
+        value = convert_rotation(value, RotationFormat.ROTMAT)
+    if is_orthonormal_rotmat(value, reduction='all'):
+        if is_rotmat_9(value):
+            assert rotation.ndim <= value.ndim + 1, f'{rotation.shape}, {value.shape}'
+            if rotation.ndim > 2:
+                rotation = expand_dims(
+                    rotation, ndim=value.ndim + 1, order=[1, -1] + [1] * (rotation.ndim - 3) + [1, 1]
+                )
+            value = rotmat_as_9(torch.matmul(rotation, rotmat_as_3x3(value)))
+        else:
+            assert rotation.ndim <= value.ndim, f'{rotation.shape}, {value.shape}'
+            if rotation.ndim > 2:
+                rotation = expand_dims(
+                    rotation, ndim=value.ndim, order=[1, -1] + [1] * (rotation.ndim - 3) + [1, 1]
+                )
+            value = torch.matmul(rotation, value)
+    else:
+        assert value.shape[-1] == 3, f'Expected a 3-dim vector in last dim, but got shape: {value.shape}'
+        assert rotation.ndim <= value.ndim + 1, f'{rotation.shape}, {value.shape}'
+        if rotation.ndim > 2:
+            rotation = expand_dims(
+                rotation, ndim=value.ndim + 1, order=[1, -1] + [1] * (rotation.ndim - 3) + [1, 1]
+            )
+        value = torch.matmul(rotation, value.unsqueeze(-1)).squeeze(-1)
+    if quaternion:
+        value = convert_rotation(value, RotationFormat.QUATERNION)
+    return value
+
+
+def relative_to_delta_rotations(
+    rotation_sequence: torch.Tensor, encoding_frame: ReferenceFrame
+) -> torch.Tensor:
+    """
+    Transform a sequence of rotation representations encoded w.r.t. the same reference frame to delta
+    rotations where each element is encoded w.r.t. the PREVIOUS rotation frame in the sequence.
+    The first element in the sequence remains the same.
+
+    Ex:
+        Sequence of points (rotations): R_1, R_2, R_3, R_4
+        `rotation_sequence` contains the rotations: R_01, R_02, R_03, R_04, where 0 is the reference frame
+            and R_01 is the pose of R1 frame in the reference frame 0, i.e. R_10 converts from reference
+            frame to R1 frame
+        Output: R_01, R_12, R_23, R_34, i.e. the rotation poses of R_1 in 0 frame, of R_2 in R1 frame, etc
+
+    Args:
+        rotation_sequence: torch.Tensor of shape [..., S, 9], [..., S, 3, 3] or [..., S, 4], containing
+            either rotation matrices (R_01, R_12, R_23, R_34, ...) or quaternions, where S corresponds
+            to the sequence dimension
+        encoding_frame: Indicates the frame w.r.t. which the input rotations are expressed.
+            - EEF: Input rotations are fully expressed w.r.t. 0-th reference frame,
+                (i.e. the axis of rotation is defined in 0-th reference frame)
+                    R_12 = R_01^-1 @ R_02
+                    R_23 = R_12^-1 @ R_03
+            - ROBOT_BASE: Input rotations are still relative, but the
+                axis of rotation is defined in robot base frame
+                    R_12 = R_01^-1 @ R_02
+                    R_23 = R_12^-1 @ R_03
+            - All other EEF or ROBOT_BASE frames treated accordingly
+    Returns:
+        torch.Tensor of the same shape as rotation_sequence, containing delta rotations
+    """
+    assert rotation_sequence.ndim >= 3, rotation_sequence.shape
+    rotation_format: RotationFormat = rotation_format_from_tensor(rotation_sequence)
+    rotation_sequence = convert_rotation(rotation_sequence, RotationFormat.QUATERNION)
+    reference_sequence = torch.roll(rotation_sequence, 1, dims=-2).clone()
+    reference_sequence[..., 0, :] = roma.identity_quat()
+    reference_sequence = roma.quat_inverse(reference_sequence)
+    if encoding_frame in ReferenceFrame.eef_frames:
+        delta_rotations = roma.quat_product(reference_sequence, rotation_sequence)
+    elif encoding_frame in ReferenceFrame.robot_frames:
+        delta_rotations = roma.quat_product(rotation_sequence, reference_sequence)
+    else:
+        raise NotImplementedError(f'Encoding frame {encoding_frame} not implemented')
+    delta_rotations = convert_rotation(delta_rotations, rotation_format)
+    return delta_rotations
+
+
+def delta_to_relative_rotations(
+    rotation_sequence: torch.Tensor, encoding_frame: ReferenceFrame
+) -> torch.Tensor:
+    """
+    Transform a sequence of rotation representations encoded w.r.t. the PREVIOUS rotation frame in the
+    sequence to the 0-th element preceding the sequence
+
+    Ex:
+        `rotation_sequence` contains the rotations: R_01, R_12, R_23, R_34, where R0 is the base frame,
+            implicitly encoded in R_01 and R_10 converts from R0 frame to R1 frame
+        Output: R_01, R_02, R_03, R_04
+
+    Args:
+        rotation_sequence: torch.Tensor of shape [..., S, 9], [..., S, 3, 3] or [..., S, 4], containing
+            either rotation matrices (R_01, R_12, R_23, R_34, ...) or quaternions
+    Returns:
+        torch.Tensor of shape [..., S, 9], [..., S, 3, 3] or [..., S, 4] containing transformed rotations
+            (R_01, R_02, R_03, R_04, ...)
+    """
+    assert rotation_sequence.ndim >= 3, rotation_sequence.shape
+    rotation_format: RotationFormat = rotation_format_from_tensor(rotation_sequence)
+    rotation_sequence = convert_rotation(rotation_sequence, RotationFormat.QUATERNION)
+    rotation_sequence = rotation_sequence.clone()
+    cumulative = rotation_sequence[..., :1, :]
+    delta_rotations = [cumulative]
+    for i in range(2, rotation_sequence.shape[-2] + 1):
+        if encoding_frame in ReferenceFrame.eef_frames:
+            cumulative = roma.quat_product(cumulative, rotation_sequence[..., i - 1 : i, :])
+        elif encoding_frame in ReferenceFrame.robot_frames:
+            cumulative = roma.quat_product(rotation_sequence[..., i - 1 : i, :], cumulative)
+        else:
+            raise NotImplementedError(f'Encoding frame {encoding_frame} not implemented')
+        delta_rotations.append(cumulative)
+    delta_rotations = torch.cat(delta_rotations, dim=-2)
+    delta_rotations = convert_rotation(delta_rotations, rotation_format)
+    return delta_rotations
+
+
+def world_to_relative_rotations(
+    rotation_sequence: torch.Tensor, reference_rotation: torch.Tensor, encoding_frame: ReferenceFrame
+) -> torch.Tensor:
+    """
+    Transform a sequence of rotations expressed w.r.t. WORLD frame to relative rotations w.r.t.
+    `reference_rotation`, where `reference_rotation` is provided w.r.t. WORLD frame.
+
+    Ex:
+        Sequence of points (rotations): R_0, R_1, R_2, R_3, R_4
+        `rotation_sequence` contains the rotations: R_W1, R_W2, R_W3, R_W4, where W is the world frame
+            and R_W1 is the pose of R1 frame in world frame, i.e. R_1W converts from world frame to R1 frame
+        `reference_rotation`: R_W0
+        Output: R_01, R_02, R_03, R_04 -> the rotation poses of R_1, R_2, R_3, R_4 expressed in R_0 frame
+
+    Args:
+        rotation_sequence: torch.Tensor of shape [..., S, 9], [..., S, 3, 3] or [..., S, 4], containing
+            either rotation matrices (R_W1, R_W2, R_W3, R_W4, ...) or quaternions
+        reference_rotation: torch.Tensor, shape [..., 9], [..., 3, 3] or [..., 4] and the SAME number of BATCH
+            dims as `rotation_sequence`. The new reference frame, provided w.r.t. WORLD coordinate frame R_W0
+        encoding_frame: Indicates the frame w.r.t. which the output rotations would be encoded - the fixed
+            world frame (ROBOT_BASE) or the local reference_frame (EEF)
+            - EEF: Output rotations are fully expressed w.r.t. reference_rotation
+                (i.e. the axis of rotation is defined in reference frame)
+                    R_W1 = R_W0 @ R_01  <=>  R_01 = R_0W @ R_W1
+            - ROBOT_BASE: Output rotations are still relative, but
+                the axis of rotation is defined in robot base frame
+                    R_W1 = R_01 @ R_W0  <=>  R_01 = R_W1 @ R_0W
+            - All other EEF or ROBOT_BASE frames treated accordingly
+    Returns:
+        torch.Tensor of shape [..., S, 9], [..., S, 3, 3] or [..., S, 4] containing transformed rotations
+            (R_01, R_02, R_03, R_04, ...)
+    """
+    assert rotation_sequence.ndim >= 3, rotation_sequence.shape
+    rotation_format: RotationFormat = rotation_format_from_tensor(rotation_sequence)
+    reference_rotation = rotmat_as_3x3(convert_rotation(reference_rotation, RotationFormat.ROTMAT))
+    rotation_sequence = rotmat_as_3x3(convert_rotation(rotation_sequence, RotationFormat.ROTMAT))
+    if reference_rotation.ndim != rotation_sequence.ndim:
+        raise ValueError(
+            f'Cannot broadcast reference_rotation of shape {reference_rotation.shape} to rotation_sequence of shape {rotation_sequence.shape}. Provide tensors with the same number of batch dimensions'
+        )
+    R_0W = rotmat_as_3x3(rotmat_inverse(reference_rotation))
+    if encoding_frame in ReferenceFrame.eef_frames:
+        relative_rotations = torch.matmul(R_0W, rotation_sequence)
+    elif encoding_frame in ReferenceFrame.robot_frames:
+        relative_rotations = torch.matmul(rotation_sequence, R_0W)
+    else:
+        raise NotImplementedError(f'Encoding frame {encoding_frame} not implemented')
+    relative_rotations = convert_rotation(relative_rotations, rotation_format)
+    return relative_rotations
+
+
+def rotation_to_target_frame(
+    rotation: torch.Tensor,
+    source_frame: ReferenceFrame,
+    target_frame: ReferenceFrame,
+    ee_pose_rotation: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Convert rotation sequence from source_frame to target_frame
+    Args:
+        rotation: torch.Tensor of shape [..., S, 9 | 4 | 3 x 3], containing
+            the rotations, where S corresponds to the sequence dimension
+        source_frame: indicates the frame w.r.t. which `rotation` is expressed
+        target_frame: indicates the frame w.r.t. which the output rotation should be expressed
+        ee_pose_rotation: torch.Tensor of shape [..., 9 | 4 | 3 x 3], containing the rotation of the
+            current end-effector pose w.r.t. ROBOT_BASE frame. Required only when source_frame and
+            target_frame have different core reference frames.
+    Returns:
+        torch.Tensor of the same shape as rotation, containing the converted rotations
+    """
+    if source_frame == target_frame:
+        return rotation
+    assert source_frame in ReferenceFrame.robot_frames | ReferenceFrame.eef_frames, source_frame
+    assert target_frame in ReferenceFrame.robot_frames | ReferenceFrame.eef_frames, target_frame
+    if ee_pose_rotation is not None:
+        ee_pose_rotation = rotmat_as_3x3(convert_rotation(ee_pose_rotation, RotationFormat.ROTMAT))
+    if source_frame.to_core() != target_frame.to_core():
+        assert ee_pose_rotation is not None, f'{source_frame}, {target_frame}'
+        if source_frame in ReferenceFrame.delta_frames:
+            rotation = delta_to_relative_rotations(rotation, encoding_frame=source_frame)
+            source_frame = source_frame.to_relative()
+        if target_frame in ReferenceFrame.robot_frames:
+            assert source_frame == ReferenceFrame.EEF_RELATIVE, source_frame
+            rotation = world_to_relative_rotations(
+                rotation, reference_rotation=rotmat_inverse(ee_pose_rotation), encoding_frame=source_frame
+            )
+            source_frame = ReferenceFrame.ROBOT_BASE
+        elif target_frame in ReferenceFrame.eef_frames:
+            assert source_frame in ReferenceFrame.robot_frames, source_frame
+            if source_frame == ReferenceFrame.ROBOT_BASE_RELATIVE:
+                rotation = world_to_relative_rotations(
+                    rotation, reference_rotation=rotmat_inverse(ee_pose_rotation), encoding_frame=source_frame
+                )
+                source_frame = ReferenceFrame.ROBOT_BASE
+            rotation = world_to_relative_rotations(
+                rotation, reference_rotation=ee_pose_rotation, encoding_frame=target_frame
+            )
+            source_frame = target_frame.to_relative()
+    assert source_frame.to_core() == target_frame.to_core(), f'{source_frame}, {target_frame}'
+    if source_frame == target_frame:
+        return rotation
+    if (
+        source_frame in ReferenceFrame.delta_frames
+        and target_frame in ReferenceFrame.relative_frames | ReferenceFrame.core_frames
+    ):
+        rotation = delta_to_relative_rotations(rotation, encoding_frame=source_frame)
+        source_frame = source_frame.to_relative()
+    elif source_frame == ReferenceFrame.ROBOT_BASE:
+        assert ee_pose_rotation is not None
+        rotation = world_to_relative_rotations(
+            rotation, reference_rotation=ee_pose_rotation, encoding_frame=source_frame
+        )
+        source_frame = ReferenceFrame.ROBOT_BASE_RELATIVE
+    assert source_frame in ReferenceFrame.relative_frames, source_frame
+    if target_frame in ReferenceFrame.delta_frames:
+        rotation = relative_to_delta_rotations(rotation, encoding_frame=source_frame)
+        source_frame = source_frame.to_delta()
+    elif target_frame == ReferenceFrame.ROBOT_BASE:
+        rotation = world_to_relative_rotations(
+            rotation, reference_rotation=rotmat_inverse(ee_pose_rotation), encoding_frame=source_frame
+        )
+        source_frame = ReferenceFrame.ROBOT_BASE
+    assert source_frame == target_frame, f'{source_frame}, {target_frame}'
+    return rotation
+
+
+def stereographic_map_quaternion(
+    quaternion: torch.Tensor, k: float, inverse: bool, eps: float = 1e-08
+) -> torch.Tensor:
+    """
+    Forward or inverse 1-1 quaternion remapping on S^3 using stereographic linear map.
+    Forward map:
+        theta' = 4 * arctan(k * tan(theta / 4)) where q = [cos(theta), sin(theta)*axis]
+    Inverse map:
+        theta = 4 * arctan( 1/k * tan(theta' / 4))
+
+    Args:
+        quaternion: torch.Tensor of shape [..., 4], input quaternion
+        k: positive scalar stretch factor
+        eps: numerical stability constant.
+
+    Returns:
+        torh.Tensor of shape [..., 4], mapped quaternion
+    """
+    assert k > 0, f'Stretch factor k must be positive, but got {k}'
+    assert is_quaternion(quaternion), f'{quaternion.shape} not a quaternion'
+    rotvec = roma.unitquat_to_rotvec(quaternion)
+    theta = torch.norm(rotvec, dim=-1, keepdim=True)
+    k_eff = k if not inverse else 1.0 / k
+    theta_prime = 4.0 * torch.atan(k_eff * torch.tan(torch.clamp(theta / 4.0, min=0, max=torch.pi / 2 - eps)))
+    rotvec = rotvec / torch.max(theta, torch.tensor(eps)) * theta_prime
+    quaternion_output = roma.rotvec_to_unitquat(rotvec)
+    return quaternion_output
+
+
+def stereographic_map_rotation(
+    rotation: torch.Tensor, factor: float, inverse: bool, eps=1e-08
+) -> torch.Tensor:
+    if factor == 1.0:
+        return rotation
+    rotation_format = rotation_format_from_tensor(rotation)
+    is_3x3 = is_rotmat_3x3(rotation)
+    rotation = convert_rotation(rotation, RotationFormat.QUATERNION, autonorm=False, half_cover=True)
+    rotation = stereographic_map_quaternion(rotation, factor, inverse=inverse, eps=eps)
+    rotation = convert_rotation(rotation, rotation_format, autonorm=False, half_cover=True)
+    if is_3x3:
+        rotation = rotmat_as_3x3(rotation)
+    return rotation
+
+
+class RotationStereomapNormalizer(Normalizer[RotationStereomapNormalizerConfig]):
+    def normalize(self, value: torch.Tensor, **kwargs) -> torch.Tensor:
+        del kwargs
+        return stereographic_map_rotation(value, factor=self.config.factor, inverse=False)
+
+    def unnormalize(self, value: torch.Tensor, **kwargs) -> torch.Tensor:
+        del kwargs
+        return stereographic_map_rotation(value, factor=self.config.factor, inverse=True)
+
+
+def assert_np_hwc_or_hw_image(image: np.ndarray | PIL.Image.Image) -> np.ndarray:
+    """Make sure image is of type np.ndarray and HWC format"""
+    if isinstance(image, PIL.Image.Image):
+        image = np.asarray(image)
+    assert isinstance(image, np.ndarray), type(image)
+    assert image.ndim in [2, 3], image.shape
+    if image.ndim == 3:
+        assert image.shape[-1] <= 4, image.shape
+    return image
+
+
+def hw_from_image(image: PIL.Image.Image | np.ndarray) -> tuple[int, int]:
+    if isinstance(image, np.ndarray):
+        (height, width) = image.shape[:2]
+    else:
+        (width, height) = image.size
+    return height, width
+
+
+def pad_image(
+    image: PIL.Image.Image | np.ndarray,
+    target_size: dict[str, int],
+    pad_value: tuple[int, int, int] | tuple[float, float, float] | int | float = 0,
+) -> PIL.Image.Image | np.ndarray:
+    """Pad image adding a symmetric border around the height/width."""
+    assert isinstance(image, (PIL.Image.Image, np.ndarray)), type(image)
+    (height, width) = hw_from_image(image)
+    (target_width, target_height) = (target_size['width'], target_size['height'])
+    if width == target_width and height == target_height:
+        return image
+    assert target_width >= width, f"Can't pad image of width {width} to {target_width}"
+    assert target_height >= height, f"Can't pad image of height {height} to {target_height}"
+    (horizontal_pad, vertical_pad) = (int((target_width - width) / 2), int((target_height - height) / 2))
+    if isinstance(image, np.ndarray):
+        padding = ((vertical_pad, vertical_pad), (horizontal_pad, horizontal_pad)) + ((0, 0),) * (
+            image.ndim - 2
+        )
+        image = np.pad(image, padding, mode='constant', constant_values=pad_value)
+    else:
+        padding = (horizontal_pad, vertical_pad, horizontal_pad, vertical_pad)
+        image = torchvision.transforms.v2.functional.pad(
+            image, padding=padding, fill=pad_value, padding_mode='constant'
+        )
+    return image
+
+
+def pad_image_to_ratio(
+    image: PIL.Image.Image | np.ndarray,
+    target_wh_ratio: float,
+    pad_value: tuple[int, int, int] | tuple[float, float, float] | int | float = 0,
+) -> PIL.Image.Image | np.ndarray:
+    """Pad image to a target aspect ratio."""
+    (height, width) = hw_from_image(image)
+    wh_ratio = width / height
+    if target_wh_ratio >= wh_ratio:
+        pad_size = {'width': round(height * target_wh_ratio), 'height': height}
+    else:
+        pad_size = {'width': width, 'height': round(width / target_wh_ratio)}
+    image = pad_image(image, target_size=pad_size, pad_value=pad_value)
+    return image
+
+
+def crop_image(
+    image: np.ndarray | PIL.Image.Image,
+    start_height: int,
+    start_width: int,
+    target_height: int,
+    target_width: int,
+) -> np.ndarray | PIL.Image.Image:
+    np_image = assert_np_hwc_or_hw_image(image)
+    (height, width) = hw_from_image(image)
+    assert target_width <= width, f"Can't crop image of width {width} to {target_width}"
+    assert target_height <= height, f"Can't crop image of width {height} to {target_height}"
+    (start_height, start_width) = (round(start_height), round(start_width))
+    (target_height, target_width) = (round(target_height), round(target_width))
+    np_image = np_image[
+        start_height : start_height + target_height, start_width : start_width + target_width, ...
+    ]
+    image = PIL.Image.fromarray(np_image) if isinstance(image, PIL.Image.Image) else np_image
+    return image
+
+
+def crop_image_center(
+    image: np.ndarray | PIL.Image.Image, target_size: dict[str, int]
+) -> np.ndarray | PIL.Image.Image:
+    np_image = assert_np_hwc_or_hw_image(image)
+    (height, width) = np_image.shape[:2]
+    (target_height, target_width) = (target_size['height'], target_size['width'])
+    assert target_width <= width, f"Can't crop image of width {width} to {target_width}"
+    assert target_height <= height, f"Can't crop image of width {height} to {target_height}"
+    top = (height - target_height) // 2
+    left = (width - target_width) // 2
+    np_image = crop_image(np_image, top, left, target_height, target_width)
+    image = PIL.Image.fromarray(np_image) if isinstance(image, PIL.Image.Image) else np_image
+    return image
+
+
+def crop_image_to_ratio(
+    image: PIL.Image.Image | np.ndarray, target_wh_ratio: float
+) -> PIL.Image.Image | np.ndarray:
+    """Pad image to a target aspect ratio."""
+    (height, width) = hw_from_image(image)
+    wh_ratio = width / height
+    if target_wh_ratio >= wh_ratio:
+        crop_size = {'width': width, 'height': round(width / target_wh_ratio)}
+    else:
+        crop_size = {'width': round(height * target_wh_ratio), 'height': height}
+    image = crop_image_center(image, target_size=crop_size)
+    return image
+
+
+def crop_and_pad_image_to_ratio(
+    image: PIL.Image.Image | np.ndarray,
+    target_wh_ratio: float,
+    mode: ResizeMode | str,
+    pad_value: tuple[int, int, int] | tuple[float, float, float] | int | float = 0,
+) -> PIL.Image.Image | np.ndarray:
+    """
+    Crop and pad an image to a target size depending on the mode.
+    It's expected that the source image and target size have different aspect ratios.
+
+    Args:
+        image: The image to crop and pad.
+        target_size: The target size to crop and pad the image to.
+        mode: The mode to use for cropping and padding.
+    """
+    (height, width) = hw_from_image(image)
+    wh_ratio = width / height
+    if np.isclose(wh_ratio, target_wh_ratio, rtol=0.01, atol=0.0001):
+        return image
+    if mode == ResizeMode.SMART:
+        aspect_ratio = max(width, height) / min(width, height)
+        target_ratio = max(target_wh_ratio, 1 / target_wh_ratio)
+        if aspect_ratio == 1:
+            if target_ratio >= 4 / 3 - 0.01:
+                crop_wh_ratio = 4 / 3 if target_wh_ratio >= 1.0 else 3 / 4
+                image = crop_image_to_ratio(image, crop_wh_ratio)
+            else:
+                pass
+        elif aspect_ratio <= 4 / 3 + 0.01:
+            if wh_ratio >= 1.0 != (target_wh_ratio >= 1.0):
+                image = crop_image_to_ratio(image, 1.0)
+        elif wh_ratio >= 1.0 != (target_wh_ratio >= 1.0):
+            image = crop_image_to_ratio(image, 1.0)
+        elif target_ratio >= 4 / 3 + 0.01:
+            pass
+        else:
+            crop_wh_ratio = 4 / 3 if target_wh_ratio >= 1.0 else 3 / 4
+            image = crop_image_to_ratio(image, crop_wh_ratio)
+        image = pad_image_to_ratio(image, target_wh_ratio, pad_value=pad_value)
+    elif mode == ResizeMode.PAD:
+        image = pad_image_to_ratio(image, target_wh_ratio, pad_value=pad_value)
+    elif mode == ResizeMode.CROP:
+        image = crop_image_to_ratio(image, target_wh_ratio)
+    else:
+        raise ValueError(f'Mode {mode} not supported')
+    return image
+
+
+def is_single_channel_image(image: np.ndarray | PIL.Image.Image) -> bool:
+    if isinstance(image, PIL.Image.Image):
+        return image.mode in ['1', 'L', 'LA', 'La', 'P', 'PA', 'F', 'I', 'I;16', 'I;16L', 'I;16B', 'I;16N']
+    if isinstance(image, np.ndarray):
+        return image.ndim == 2 or image.ndim == 3 and image.shape[2] == 1
+    raise ValueError(f'Unsupported image type: {type(image)}')
+
+
+def is_binary_mask(image: np.ndarray | PIL.Image.Image) -> bool:
+    image = np.asarray(image)
+    return image.dtype in [np.uint8, np.bool_] and np.max(image) == 1
+
+
+def resize_image(
+    image: PIL.Image.Image | np.ndarray,
+    target_size: dict[str, int],
+    mode: ResizeMode | str,
+    resample: PIL.Image.Resampling | str = 'auto',
+    pad_value: tuple[int, int, int] | tuple[float, float, float] | int | float = 0,
+) -> PIL.Image.Image | np.ndarray:
+    (target_width, target_height) = (target_size['width'], target_size['height'])
+    (height, width) = hw_from_image(image)
+    if height == target_height and width == target_width:
+        return image
+    if resample == 'auto':
+        if is_single_channel_image(image):
+            resample = PIL.Image.Resampling.BILINEAR
+        else:
+            resample = PIL.Image.Resampling.LANCZOS
+    else:
+        assert isinstance(resample, PIL.Image.Resampling), resample
+        if is_single_channel_image(image) and resample not in [
+            PIL.Image.Resampling.BILINEAR,
+            PIL.Image.Resampling.BICUBIC,
+        ]:
+            raise ValueError(
+                f'Single channel images must be resized with bilinear or bicubic, but got {resample}'
+            )
+    if is_bin_mask := is_binary_mask(image):
+        image = np.asarray(image).astype(np.uint8) * 255
+    if mode == ResizeMode.SMART:
+        image = crop_and_pad_image_to_ratio(
+            image, target_wh_ratio=target_width / target_height, mode=mode, pad_value=pad_value
+        )
+    pil_image = PIL.Image.fromarray(image) if isinstance(image, np.ndarray) else image
+    if mode in [ResizeMode.NAIVE, ResizeMode.SMART]:
+        pil_image = pil_image.resize((target_width, target_height), resample=resample)
+    else:
+        raise NotImplementedError(f'Mode {mode} not supported')
+    image = np.asarray(pil_image) if isinstance(image, np.ndarray) else pil_image
+    if is_bin_mask:
+        image = image.astype(np.uint8) > 127
+    return image
+
+
+def invert_gripper(gripper: np.ndarray, low: float, high: float) -> np.ndarray:
+    if low < 0.0:
+        return np.clip(-gripper, low, high)
+    return high - np.clip(gripper, low, high)
+
+
+GRIPPER_BOUNDS = {
+    'austin_buds_dataset': (0.0, 0.08),
+    'austin_sailor_dataset': (0.0, 0.08),
+    'austin_sirius_dataset': (0.0, 0.08),
+    'bc_z': (0.0, 1.0),
+    'berkeley_autolab_ur5': (0.0, 1.0),
+    'berkeley_cable_routing': (0.0, 1.0),
+    'berkeley_fanuc_manipulation': (0.0, 1.0),
+    'bridge': (0.0, 1.0),
+    'bridge_steering': (0.0, 1.0),
+    'bridge_nils': (0.0, 1.0),
+    'bridge_tread': (0.0, 1.0),
+    'bridge_paraphrase_k5': (0.0, 1.0),
+    'bridge_paraphrase_k10': (0.0, 1.0),
+    'bridge_full_tread_8b_k5': (0.0, 1.0),
+    'bridge_tread_full': (0.0, 1.0),
+    'bridge_coarse_max3': (0.0, 1.0),
+    'bridge_hindsight': (0.0, 1.0),
+    'bridge_32b': (0.0, 1.0),
+    'bridge_tread_k10': (0.0, 1.0),
+    'bridge_paraphrase_k5_mix50': (0.0, 1.0),
+    'bridge_rich_properties': (0.0, 1.0),
+    'bridge_rich_properties_full': (0.0, 1.0),
+    'bridge_rich_properties_p30': (0.0, 1.0),
+    'bridge_rich_properties_p50': (0.0, 1.0),
+    'bridge_rich_properties_mix50': (0.0, 1.0),
+    'bridge_orig': (0.0, 1.0),
+    'cmu_stretch': (-3.0, 3.0),
+    'dlr_edan_shared_control': (0.0, 1.0),
+    'droid': (0.0, 1.0),
+    'fmb': (0.0, 1.0),
+    'fractal20220817_data': (0.0, 1.0),
+    'furniture_bench_dataset': (0.0, 0.08),
+    'iamlab_cmu_pickup_insert': (0.0, 1.0),
+    'jaco_play': (0.0, 1.4),
+    'kuka': (0.0, 1.0),
+    'language_table': (0.0, 1.0),
+    'nyu_franka_play_dataset': (0.0, 1.0),
+    'roboset': (0.0, 1.0),
+    'roboturk': (0.0, 1.0),
+    'stanford_hydra_dataset': (0.0, 0.08),
+    'taco_play': (0.0, 0.08),
+    'toto': (0.0, 1.0),
+    'ucsd_kitchen_dataset': (0.0, 1.0),
+    'utaustin_mutex': (0.0, 0.08),
+    'viola': (0.0, 0.08),
+}
+
+
+def preprocess_gripper_observation(
+    gripper: np.ndarray, dataset_name: str | np.ndarray, binary: bool = True
+) -> np.ndarray:
+    """
+    Preprocess gripper observation depending on dataset. Input is the raw gripper observation from the dataset
+    or from the robot and output is normalized continuous value.
+        - if `binary`, output is in [0, 1], with 0 = closed and 1 = open.
+        - otherwise, output is in [-1, 1], with -1 = closed and 1 = open.
+
+    Dataset-specific gripper observations:
+        austin_buds_dataset: continuous; ~[0=closed; 0.08=open] (franka gripper)
+        austin_sailor_dataset: continuous; ~[0=closed; 0.08=open] (franka gripper)
+        austin_sirius_dataset: continuous; ~[0=closed; 0.08=open] (franka gripper)
+        bc_z: continuous; [0=open; 1=closed]
+        berkeley_autolab_ur5: binary; [0=open; 1=closed]
+        berkeley_cable_routing: constant (closed)
+        berkeley_fanuc_manipulation: binary; [0=open; 1=closed]
+        bridge: continuous; ~[0=closed; 1=open]
+        bridge_orig: continuous; ~[0=closed; 1=open]
+        cmu_stretch: continuous; [-3=closed; 3=open]
+        dlr_edan_shared_control: missing
+        droid: continuous; [0=open, 1=closed]
+        fmb: binary; [0=open; 1=closed]
+        fractal20220817_data: continuous; [0=open; 1=closed]
+        furniture_bench_dataset: continuous; ~[0=closed; 0.08=open] (franka gripper)
+        iamlab_cmu_pickup_insert: binary; [0=closed; 1=open]
+        jaco_play: continuous; [0=open; 1.4=closed]
+        kuka: binary; [0=open; 1=closed]
+        language_table: constant (no gripper)
+        nyu_franka_play_dataset: missing
+        roboset: continuous; [0=open, 1=closed]
+        roboturk: continuous; [0=closed, 0.04=open]
+        stanford_hydra_dataset: continuous; ~[0=closed; 0.08=open] (franka gripper)
+        taco_play: continuous; ~[0=closed; 0.08=open] (franka gripper)
+        toto: constant (closed)
+        ucsd_kitchen_dataset: missing
+        utaustin_mutex: continuous; ~[0=closed; 0.08=open] (franka gripper)
+        viola: continuous; ~[0=closed; 0.08=open] (franka gripper)
+
+    """
+    if isinstance(dataset_name, np.ndarray):
+        assert np.unique(dataset_name).size == 1, dataset_name
+        dataset_name = str(dataset_name[0])
+    if dataset_name in [
+        'berkeley_cable_routing',
+        'dlr_edan_shared_control',
+        'language_table',
+        'nyu_franka_play_dataset',
+        'toto',
+        'ucsd_kitchen_dataset',
+    ]:
+        gripper = normalize_gripper_by_bounds(
+            torch.from_numpy(gripper),
+            low=torch.full(gripper.shape, GRIPPER_BOUNDS[dataset_name][0], dtype=torch.float32),
+            high=torch.full(gripper.shape, GRIPPER_BOUNDS[dataset_name][1], dtype=torch.float32),
+            binary=binary,
+        ).numpy()
+    elif dataset_name in [
+        'bc_z',
+        'berkeley_autolab_ur5',
+        'berkeley_fanuc_manipulation',
+        'droid',
+        'fmb',
+        'fractal20220817_data',
+        'jaco_play',
+        'kuka',
+        'roboset',
+    ]:
+        (low, high) = GRIPPER_BOUNDS[dataset_name]
+        gripper = normalize_gripper_by_bounds(
+            torch.from_numpy(invert_gripper(gripper, low=low, high=high)),
+            low=torch.full(gripper.shape, GRIPPER_BOUNDS[dataset_name][0], dtype=torch.float32),
+            high=torch.full(gripper.shape, GRIPPER_BOUNDS[dataset_name][1], dtype=torch.float32),
+            binary=binary,
+        ).numpy()
+    elif dataset_name in [
+        'austin_buds_dataset',
+        'austin_sailor_dataset',
+        'austin_sirius_dataset',
+        'bridge',
+        'bridge_steering',
+        'bridge_nils',
+        'bridge_tread',
+        'bridge_paraphrase_k5',
+        'bridge_paraphrase_k10',
+        'bridge_full_tread_8b_k5',
+        'bridge_tread_full',
+        'bridge_coarse_max3',
+        'bridge_hindsight',
+        'bridge_32b',
+        'bridge_tread_k10',
+        'bridge_paraphrase_k5_mix50',
+        'bridge_rich_properties',
+        'bridge_rich_properties_full',
+        'bridge_rich_properties_p30',
+        'bridge_rich_properties_p50',
+        'bridge_rich_properties_mix50',
+        'bridge_orig',
+        'cmu_stretch',
+        'furniture_bench_dataset',
+        'iamlab_cmu_pickup_insert',
+        'roboturk',
+        'stanford_hydra_dataset',
+        'taco_play',
+        'utaustin_mutex',
+        'viola',
+    ]:
+        (low, high) = GRIPPER_BOUNDS[dataset_name]
+        gripper = normalize_gripper_by_bounds(
+            torch.from_numpy(gripper),
+            low=torch.full(gripper.shape, low, dtype=torch.float32),
+            high=torch.full(gripper.shape, high, dtype=torch.float32),
+            binary=binary,
+        ).numpy()
+    else:
+        raise NotImplementedError(f'Unknown dataset: {dataset_name}')
+    return gripper
+
+
+VLMProcessorConfigT = TypeVar('VLMProcessorConfigT', bound=VLMProcessorConfig)
+
+
+class VLMProcessor(Configurable[VLMProcessorConfigT], Template[VLMProcessorConfigT]):
+    @abstractmethod
+    def preprocess_inputs(
+        self, chat: List[str], images: Dict[str, List[PIL.Image.Image]]
+    ) -> Dict[str, torch.Tensor | Dict[str, torch.Tensor]]:
+        ...
+
+    @property
+    @abstractmethod
+    def tokenizer(self) -> transformers.PreTrainedTokenizerBase:
+        pass
+
+    @property
+    @abstractmethod
+    def image_sizes(self) -> Dict[str, ImageSizeConfig]:
+        pass
+
+    @property
+    @abstractmethod
+    def ignore_index(self) -> int:
+        pass
+
+
+VLAMProcessorConfigT = TypeVar('VLAMProcessorConfigT', bound=VLAMProcessorConfig)
+
+
+class VLAMProcessor(Configurable[VLAMProcessorConfigT], Template[VLAMProcessorConfigT]):
+    def __init__(self, config: VLAMProcessorConfigT, vlm_processor: VLMProcessor):
+        super().__init__(config)
+        self.vlm_processor = vlm_processor
+        self.control_tokenizer = EmptyTokenizer(
+            config=self.config.control_tokenizer_config, tokenizer=self.tokenizer
+        )
+        self.translation_obs_norm = DatasetStatsNormalizer(self.config.translation_obs_norm)
+        self.rotation_obs_norm = IdentityNormalizer(self.config.rotation_obs_norm)
+        self.translation_control_norm = BoundsNormalizer(self.config.translation_control_norm)
+        self.rotation_control_norm = RotationStereomapNormalizer(self.config.rotation_control_norm)
+        self.joints_obs_norm = BoundsNormalizer(self.config.joints_obs_norm)
+
+    @property
+    def tokenizer(self) -> transformers.PreTrainedTokenizerBase:
+        return self.vlm_processor.tokenizer
+
+    @property
+    def image_sizes(self) -> Dict[str, ImageSizeConfig]:
+        return self.vlm_processor.image_sizes
+
+    @property
+    def camera_names(self) -> List[str]:
+        return list(self.vlm_processor.image_sizes.keys())
+
+    @property
+    def ignore_index(self) -> int:
+        return self.vlm_processor.ignore_index
+
+    @property
+    def control_io_config(self) -> ControlDataIOConfig:
+        return self.config.control_io_config
+
+    @cached_property
+    def rotation_components(self) -> int:
+        if self.config.rotation_format == RotationFormat.EULER:
+            return 3
+        if self.config.rotation_format == RotationFormat.QUATERNION:
+            return 4
+        if self.config.rotation_format == RotationFormat.ROTMAT:
+            return 9
+        raise NotImplementedError(self.config.rotation_format)
+
+    @abstractmethod
+    def policy_control_plan_from_model_target(
+        self, target: RoboticsTarget, dataset_name: np.ndarray
+    ) -> RoboticsControlPlan:
+        """
+        Produce a RoboticsControlPlan from `model_output`. Unnormalizes the outputs, runs any
+        model-specific postprocessing and converts to the desired target reference frame.
+        See `policy_control_plan_from_model_output` for details on arguments.
+        """
+
+    @abstractmethod
+    def policy_control_plan_from_model_output(
+        self, model_output: RoboticsOutput, dataset_name: np.ndarray, valid_mask: torch.Tensor
+    ) -> RoboticsControlPlan:
+        """
+        Produce a RoboticsControlPlan from `model_output`. Unnormalizes the outputs and runs any
+        model-specific postprocessing. Translation and rotation outputs are always in a RELATIVE
+        frame w.r.t. the currrent end-effector pose, where the reference frame used during learning
+        (ROBOT_BASE vs EEF) is preserved for each component. In other words, if translation_control_frame
+        is ROBOT_BASE_DELTA, and rotation_control_frame is EEF_DELTA, then the output translation will be
+        in ROBOT_BASE_RELATIVE frame and rotation in EEF_RELATIVE frame.
+
+        We explicitly avoid any conversions which require the EE pose. The EE pose needs to be in
+        ROBOT_BASE frame, but there are many easy sources of error. For example, it's easy to mistakenly
+        provide the EE pose, which was input to the model and is not guaranteed to be in ROBOT_BASE.
+        It's also easy to provide normalized EE pose, which also leads to incorrect results. Instead,
+        if further conversions are required, it's recommended to call translation_to_target_frame and
+        rotation_to_target_frame outside this function, where the user has full control over.
+
+        Args:
+            model_output: RoboticsOutput from the model of shape [B, num_timesteps, ...]
+            dataset_name: np.ndarray of shape [B] with dataset names for each batch example
+            valid_mask: torch.Tensor of shape [B, num_timesteps] indicating valid control steps
+        Returns:
+            RoboticsControlPlan with **UNNORMALIZED** controls in the desired target frame
+        """
+
+    def resize_image(
+        self, camera_name: str, image: PIL.Image.Image | np.ndarray
+    ) -> PIL.Image.Image | np.ndarray:
+        return resize_image(
+            image,
+            target_size={
+                'width': self.image_sizes[camera_name].width,
+                'height': self.image_sizes[camera_name].height,
+            },
+            mode=self.config.image_resize,
+            resample=PIL.Image.Resampling.LANCZOS,
+        )
+
+    def preprocess_inputs(
+        self,
+        chat: List[str],
+        images: Dict[str, PIL.Image.Image | List[PIL.Image.Image]],
+        ee_pose_translation: np.ndarray,
+        ee_pose_rotation: np.ndarray,
+        gripper: np.ndarray,
+        joints: np.ndarray,
+        dataset_name: np.ndarray,
+        inference_mode: bool,
+        control_target: Optional[RoboticsTarget] = None,
+    ) -> Dict[str, torch.Tensor | Dict[str, torch.Tensor]]:
+        """
+        Preprocess the inputs for a single example
+        Args:
+            instruction: Language instruction
+            images: History of input images with increasing timestamps
+            ee_pose_translation: np.ndarray, shape [..., num_past_scalars, 3]
+            ee_pose_rotation: np.ndarray, shape [..., num_past_scalars, 3 | 4 | 9]
+            joints: np.ndarray, shape  [..., num_past_scalars, <= 7]
+            dataset_name: 1D np.ndarray
+            inference_mode: If True, prepare the input for inference (e.g. don't include target
+                any tokens in the input if relevant). If control_target is available, it should
+                still be preprocessed for test dataset comparison
+            control_target: RoboticsTarget, each component of shape
+                [..., num_control_steps, num_control_components]. Provided only when available, usually
+                during training and dataset test
+        Returns:
+            Dict containing torch.Tensor with inputs
+        """
+        del control_target, inference_mode
+        inputs = self.vlm_processor.preprocess_inputs(chat=chat, images=images)
+        images: Dict[str, torch.Tensor] = inputs['images']
+        input_ids: torch.Tensor = inputs['input_ids'][..., : self.tokenizer.model_max_length]
+        target_text_tokens_ids: torch.Tensor = inputs['target_ids'][..., : self.tokenizer.model_max_length]
+        attn_mask = torch.ones(input_ids.shape, dtype=torch.bool)
+        ee_pose_translation = torch.tensor(ee_pose_translation, dtype=torch.float32)
+        ee_pose_rotation = torch.tensor(ee_pose_rotation, dtype=torch.float32)
+        ee_pose_rotation = convert_rotation(ee_pose_rotation, self.config.rotation_format, autonorm=True)
+        gripper = preprocess_gripper_observation(gripper, dataset_name)
+        gripper = torch.tensor(gripper, dtype=torch.float32)
+        ee_pose_translation = self.normalize(
+            ee_pose_translation, dataset_name=dataset_name, key='translation_obs'
+        )
+        ee_pose_rotation = self.normalize(ee_pose_rotation, dataset_name=dataset_name, key='rotation_obs')
+        joints = torch.tensor(joints, dtype=torch.float32)
+        if joints.shape[-1] < 7:
+            missing_size = 7 - joints.shape[-1]
+            joints = torch.cat([joints, torch.zeros([*joints.shape[:-1], missing_size])], dim=-1)
+        joints = self.normalize(joints, dataset_name=dataset_name, key='joints_obs')
+        outputs = {
+            'images': images,
+            'input_ids': input_ids,
+            'target_text_tokens_ids': target_text_tokens_ids,
+            'attn_mask': attn_mask,
+            'ee_pose_translation': ee_pose_translation,
+            'ee_pose_rotation': ee_pose_rotation,
+            'gripper': gripper,
+            'joints': joints,
+            'control_tokens_ids': None,
+            'target_control_tokens_ids': None,
+        }
+        return outputs
+
+    def create_input(
+        self,
+        chat: List[str],
+        images: Dict[str, List[PIL.Image.Image]],
+        ee_pose_translation: np.ndarray,
+        ee_pose_rotation: np.ndarray,
+        gripper: np.ndarray,
+        joints: np.ndarray,
+        dataset_name: np.ndarray,
+        inference_mode: bool,
+        control_target: Optional[RoboticsTarget] = None,
+    ) -> RoboticsInput:
+        inputs = self.preprocess_inputs(
+            chat=chat,
+            images=images,
+            ee_pose_translation=ee_pose_translation,
+            ee_pose_rotation=ee_pose_rotation,
+            gripper=gripper,
+            joints=joints,
+            dataset_name=dataset_name,
+            inference_mode=inference_mode,
+            control_target=control_target,
+        )
+        inputs.pop('target_text_tokens_ids')
+        inputs.pop('target_control_tokens_ids')
+        return RoboticsInput(**inputs)
+
+    def normalize(self, value: torch.Tensor, dataset_name: np.ndarray, key: str) -> torch.Tensor:
+        normalizer = getattr(self, f'{key}_norm')
+        return normalizer.normalize(value, dataset_name=dataset_name)
+
+    def unnormalize(self, value: torch.Tensor, dataset_name: np.ndarray, key: str) -> torch.Tensor:
+        normalizer = getattr(self, f'{key}_norm')
+        return normalizer.unnormalize(value, dataset_name=dataset_name)
+
+    @property
+    def _stats_horizon_key(self) -> str:
+        if self.config.delta_controls:
+            if self.control_io_config.future_controls_sequence_stride_sec is None:
+                horizon = 0.0
+            else:
+                horizon = self.control_io_config.future_controls_sequence_stride_sec
+        elif self.control_io_config.future_controls_sequence_stride_sec is None:
+            if self.control_io_config.future_controls_sequence_length == 1:
+                horizon = 0.0
+            else:
+                raise NotImplementedError()
+        else:
+            horizon = (
+                self.control_io_config.future_controls_sequence_length
+                * self.control_io_config.future_controls_sequence_stride_sec
+            )
+        key = f'horizon_{round(horizon, 2)}s'
+        return key
+
+
+def world_to_relative_translations(
+    translation_sequence: torch.Tensor, reference_frame: torch.Tensor
+) -> torch.Tensor:
+    """
+    Transform a sequence of translation vectors encoded w.r.t. WORLD frame to encoding w.r.t.
+    `reference_frame`, where `reference_frame` is provided w.r.t. WORLD frame
+    Ex:
+        Sequence of points: T1, T2, T3, T4
+        `translation_sequence` contains the vectors: WT1, WT2, WT3, WT4, where W is the world frame
+        Output: T0T1, T0T2, T0T3, T0T4, where T0 is the reference frame
+
+    Args:
+        translation_sequence: torch.Tensor of shape [..., S, 3], containing the translation vectors, where S
+            corresponds to the sequence dimension
+        reference_frame: torch.Tensor, shape [..., 1, 3] and the SAME number of BATCH dims as
+            `translation_sequence`. The new reference frame, provided w.r.t. WORLD coordinate frame
+    Returns:
+        torch.Tensor of the same shape as translation_sequence, containing delta translations
+    """
+    assert translation_sequence.ndim >= 3, translation_sequence.shape
+    if reference_frame.ndim != translation_sequence.ndim:
+        raise ValueError(
+            f'Cannot broadcast reference_frame of shape {reference_frame.shape} to translation_sequence of shape {translation_sequence.shape}. Provide tensors with the same number of batch dimensions'
+        )
+    delta_translations = translation_sequence - reference_frame
+    return delta_translations
+
+
+def delta_to_relative_translations(translation_sequence: torch.Tensor) -> torch.Tensor:
+    """
+    Transform a sequence of translation vectors encoded w.r.t. PREVIOUS frame in the sequence to encoding
+    w.r.t. the 0-th element preceding the sequence
+    Ex:
+        Sequence of points: T1, T2, T3, T4
+        `translation_sequence` contains the vectors: T0T1, T1T2, T2T3, T3T4, where T0 is the base frame,
+        implicitly encoded in T0T1
+        Output: T0T1, T0T2, T0T3, T0T4
+
+    Args:
+        translation_sequence: torch.Tensor of shape [..., S, 3], containing the translation vectors, where S
+            corresponds to the sequence dimension
+    Returns:
+        torch.Tensor of the same shape as translation_sequence, containing delta translations
+    """
+    assert translation_sequence.ndim >= 3, translation_sequence.shape
+    delta_translations = torch.cumsum(translation_sequence, dim=-2)
+    return delta_translations
+
+
+def relative_to_delta_translations(translation_sequence: torch.Tensor) -> torch.Tensor:
+    """
+    Transform a sequence of translation vectors encoded w.r.t. the same reference frame to delta translation
+    vectors where each value is encoded w.r.t. the PREVIOUS frame in the sequence. The first element in
+    the sequence remains the same.
+    Ex:
+        Sequence of points: T1, T2, T3, T4
+        `translation_sequence` contains the vectors: RT1, RT2, RT3, RT4, where R is the reference frame
+        Output: RT1, T1T2, T2T3, T3T4
+
+    Args:
+        translation_sequence: torch.Tensor of shape [..., S, 3], containing the translation vectors, where S
+            corresponds to the sequence dimension
+    Returns:
+        torch.Tensor of the same shape as translation_sequence, containing delta translations
+    """
+    assert translation_sequence.ndim >= 3, translation_sequence.shape
+    reference_frames = torch.roll(translation_sequence, 1, dims=-2).clone()
+    reference_frames[..., 0, :] = 0
+    delta_translations = translation_sequence - reference_frames
+    return delta_translations
+
+
+def translation_to_target_frame(
+    translation: torch.Tensor,
+    source_frame: ReferenceFrame,
+    target_frame: ReferenceFrame,
+    ee_pose_translation: Optional[torch.Tensor] = None,
+    ee_pose_rotation: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Convert translation sequence from source_frame to target_frame
+    Args:
+        translation: torch.Tensor of shape [..., S, 3], containing the translation vectors, where S
+            corresponds to the sequence dimension
+        source_frame: indicates the frame w.r.t. which `translation` is expressed
+        target_frame: indicates the frame w.r.t. which the output translation should be expressed
+        ee_pose_translation: torch.Tensor of shape [B, ..., 3], containing the translation of the current
+            end-effector pose. Required only if target_frame is ROBOT_BASE and source_frame isn't.
+        ee_pose_rotation: torch.Tensor of shape [..., 9 | 4 | 3 x 3], containing the rotation of the
+            current end-effector pose w.r.t. ROBOT_BASE frame. Required only when source_frame and
+            target_frame have different core reference frames.
+    Returns:
+        torch.Tensor of the same shape as translation, containing the converted translations
+    """
+    if source_frame == target_frame:
+        return translation
+    assert source_frame in ReferenceFrame.robot_frames | ReferenceFrame.eef_frames, source_frame
+    assert target_frame in ReferenceFrame.robot_frames | ReferenceFrame.eef_frames, target_frame
+    if ee_pose_rotation is not None:
+        ee_pose_rotation = rotmat_as_3x3(convert_rotation(ee_pose_rotation, RotationFormat.ROTMAT))
+    if source_frame.to_core() != target_frame.to_core():
+        assert ee_pose_rotation is not None, f'{source_frame}, {target_frame}'
+        if source_frame in ReferenceFrame.delta_frames:
+            translation = delta_to_relative_translations(translation)
+            source_frame = source_frame.to_relative()
+        if target_frame in ReferenceFrame.robot_frames:
+            assert source_frame == ReferenceFrame.EEF_RELATIVE, source_frame
+            translation = apply_rotation(rotation=ee_pose_rotation, value=translation)
+            source_frame = ReferenceFrame.ROBOT_BASE_RELATIVE
+        elif target_frame in ReferenceFrame.eef_frames:
+            assert source_frame in ReferenceFrame.robot_frames, source_frame
+            if source_frame == ReferenceFrame.ROBOT_BASE:
+                assert ee_pose_translation is not None
+                translation = world_to_relative_translations(translation, reference_frame=ee_pose_translation)
+                source_frame = ReferenceFrame.ROBOT_BASE_RELATIVE
+            assert source_frame in ReferenceFrame.relative_frames, source_frame
+            translation = apply_rotation(rotation=rotmat_inverse(ee_pose_rotation), value=translation)
+            source_frame = ReferenceFrame.EEF_RELATIVE
+    assert source_frame.to_core() == target_frame.to_core(), f'{source_frame}, {target_frame}'
+    if source_frame == target_frame:
+        return translation
+    if (
+        source_frame in ReferenceFrame.delta_frames
+        and target_frame in ReferenceFrame.relative_frames | ReferenceFrame.core_frames
+    ):
+        translation = delta_to_relative_translations(translation)
+        source_frame = source_frame.to_relative()
+    elif source_frame == ReferenceFrame.ROBOT_BASE:
+        assert ee_pose_translation is not None
+        translation = world_to_relative_translations(translation, reference_frame=ee_pose_translation)
+        source_frame = ReferenceFrame.ROBOT_BASE_RELATIVE
+    assert source_frame in ReferenceFrame.relative_frames, source_frame
+    if target_frame in ReferenceFrame.delta_frames:
+        translation = relative_to_delta_translations(translation)
+        source_frame = source_frame.to_delta()
+    elif target_frame == ReferenceFrame.ROBOT_BASE:
+        translation = world_to_relative_translations(translation, reference_frame=-ee_pose_translation)
+        source_frame = ReferenceFrame.ROBOT_BASE
+    assert source_frame == target_frame, f'{source_frame}, {target_frame}'
+    return translation
+
+
+class RegressionProcessor(VLAMProcessor[RegressionProcessorConfig]):
+    def policy_control_plan_from_model_target(
+        self, target: RoboticsTarget, dataset_name: np.ndarray
+    ) -> RoboticsControlPlan:
+        """See VLAMProcessor.policy_control_plan_from_model_target for arguments"""
+        translation_m = self.unnormalize(
+            target.translation, dataset_name=dataset_name, key='translation_control'
+        )
+        rotation = self.unnormalize(target.rotation, dataset_name=dataset_name, key='rotation_control')
+        rotmat = convert_rotation(rotation, RotationFormat.ROTMAT)
+        gripper_prob = target.gripper
+        if self.config.translation_control_frame != ReferenceFrame.ROBOT_BASE:
+            translation_m = translation_to_target_frame(
+                translation_m,
+                source_frame=self.config.translation_control_frame,
+                target_frame=self.config.translation_control_frame.to_relative(),
+            )
+        if self.config.rotation_control_frame != ReferenceFrame.ROBOT_BASE:
+            rotmat = rotation_to_target_frame(
+                rotmat,
+                source_frame=self.config.rotation_control_frame,
+                target_frame=self.config.rotation_control_frame.to_relative(),
+            )
+        return RoboticsControlPlan(
+            translation_m=translation_m,
+            rotmat=rotmat,
+            gripper_prob=gripper_prob,
+            valid_mask=target.valid_mask,
+        )
+
+    def policy_control_plan_from_model_output(
+        self, model_output: RoboticsOutput, dataset_name: np.ndarray, valid_mask: torch.Tensor
+    ) -> RoboticsControlPlan:
+        """
+        Called during inference to create control plan from model output
+        See VLAMProcessor.policy_control_plan_from_model_output for arguments
+        """
+        translation_m = self.unnormalize(
+            model_output.translation, dataset_name=dataset_name, key='translation_control'
+        )
+        rotation = self.unnormalize(model_output.rotation, dataset_name=dataset_name, key='rotation_control')
+        rotmat = convert_rotation(rotation, RotationFormat.ROTMAT, autonorm=True)
+        gripper_prob = torch.sigmoid(model_output.gripper)
+        if self.config.translation_control_frame != ReferenceFrame.ROBOT_BASE:
+            translation_m = translation_to_target_frame(
+                translation_m,
+                source_frame=self.config.translation_control_frame,
+                target_frame=self.config.translation_control_frame.to_relative(),
+            )
+        if self.config.rotation_control_frame != ReferenceFrame.ROBOT_BASE:
+            rotmat = rotation_to_target_frame(
+                rotmat,
+                source_frame=self.config.rotation_control_frame,
+                target_frame=self.config.rotation_control_frame.to_relative(),
+            )
+        return RoboticsControlPlan(
+            translation_m=translation_m, rotmat=rotmat, gripper_prob=gripper_prob, valid_mask=valid_mask
+        )
+
+
+class PiZeroFlowMatchingProcessor(Configurable[PiZeroFlowProcessorConfig], RegressionProcessor):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.generator: torch.Generator = torch.Generator()
+
+    @cached_property
+    def beta_distribution(self) -> torch.distributions.Beta:
+        return torch.distributions.Beta(
+            self.config.distribution_hyperparams.get('alpha', 1.5),
+            self.config.distribution_hyperparams.get('beta', 1.0),
+        )
+
+    def create_input(self, *args, **kwargs) -> RoboticsFlowInput:
+        """In practice used only during inference"""
+        inputs = super().create_input(*args, **kwargs)
+        flow_input: FlowInput = self.sample_t0_input(batch_size=1, device=torch.device('cpu'))
+        inputs = RoboticsFlowInput(**inputs.as_json(), flow_input=flow_input[0, ...])
+        return inputs
+
+    def sample_timestep(self, batch_size: int) -> torch.Tensor:
+        if self.config.timestep_distribution.lower() == 'uniform':
+            eps = 1e-05
+            sample = (torch.rand(1, generator=self.generator) + torch.arange(batch_size) / batch_size) % (
+                1 - eps
+            )
+        elif self.config.timestep_distribution.lower() == 'beta':
+            sample = self.beta_distribution.sample([batch_size, 1, 1])
+            sample = (1 - self.config.sig_min) * (1 - sample)
+        else:
+            raise NotImplementedError(self.config.timestep_distribution)
+        sample = sample.view(batch_size, 1, 1)
+        return sample
+
+    def _psi_t(self, timestep: torch.Tensor, x_0: torch.Tensor, x_1: torch.Tensor) -> torch.Tensor:
+        return (1 - (1 - self.config.sig_min) * timestep) * x_0 + timestep * x_1
+
+    def _dpsi_dt(self, x_0: torch.Tensor, x_1: torch.Tensor) -> torch.Tensor:
+        return x_1 - (1 - self.config.sig_min) * x_0
+
+    def sample_t0_input(self, batch_size: int, device: torch.device) -> FlowInput:
+        if self.config.r0_distribution == 'normal':
+            controls_t0 = torch.randn(
+                [
+                    batch_size,
+                    self.config.control_io_config.future_controls_sequence_length,
+                    3 + self.rotation_components + 1,
+                ],
+                generator=self.generator,
+            ).to(device=device)
+            (translation_t0, rotation_t0, gripper_t0) = torch.split(
+                controls_t0, [3, self.rotation_components, 1], dim=-1
+            )
+            rotation_t0 = normalize_rotation(rotation_t0)
+        elif self.config.r0_distribution == 'uniform':
+            controls_t0 = torch.randn(
+                [batch_size, self.config.control_io_config.future_controls_sequence_length, 4],
+                generator=self.generator,
+            ).to(device=device)
+            (translation_t0, gripper_t0) = torch.split(controls_t0, [3, 1], dim=-1)
+            rotation_t0 = convert_rotation(
+                roma.random_unitquat(
+                    (batch_size, self.config.control_io_config.future_controls_sequence_length), device=device
+                ),
+                self.config.rotation_format,
+            )
+        else:
+            raise NotImplementedError(self.config.r0_distribution)
+        if self.config.rotation_format == RotationFormat.QUATERNION:
+            rotation_t0 = quaternion_half_cover(rotation_t0)
+        timestep = torch.zeros([batch_size, 1, 1], device=device)
+        return FlowInput(
+            timestep=timestep,
+            translation_t0=translation_t0,
+            rotation_t0=rotation_t0,
+            gripper_t0=gripper_t0,
+            translation_t=None,
+            rotation_t=None,
+            gripper_t=None,
+        )
+
+    def policy_control_plan_from_model_output(
+        self, model_output: RoboticsOutput, dataset_name: np.ndarray, valid_mask: torch.Tensor
+    ) -> RoboticsControlPlan:
+        """
+        Called during inference to create control plan from model output
+        See VLAMProcessor.policy_control_plan_from_model_output for arguments
+        """
+        model_output = model_output.replace(
+            translation=torch.clamp(model_output.translation, -1, 1),
+            rotation=torch.clamp(model_output.rotation, -1, 1),
+        )
+        control_plan = super().policy_control_plan_from_model_output(
+            model_output=model_output, dataset_name=dataset_name, valid_mask=valid_mask
+        )
+        control_plan = control_plan.replace(gripper_prob=torch.clamp(model_output.gripper, 0, 1))
+        return control_plan

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/model_config.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a16160e75e679c42a863ee98c8b3010baffca7473b30f213aef37befdb993082
+size 4124

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/raw_config.yaml

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:baa72befc2efb8a163e2615d4f733eddb33b241fbdefb667cdba10a9afaa1b72
+size 876

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/session.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7fcc08747ace4c3dcdc3a52f706c30d023be548325f7bde1f1f24f4095dc385f
+size 127896

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/session_info.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66f32905c4b9497bda291a041b76d6f9d2aa58a9abf74cd3e0aee56385021561
+size 1015

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/session_rank_0.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7fcc08747ace4c3dcdc3a52f706c30d023be548325f7bde1f1f24f4095dc385f
+size 127896

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/session_rank_1.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b6ac54e1c4fd7680858ef598bd94a590b4a8b2d37e76350bfc122f6d3ec071e
+size 4522

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/session_rank_2.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66c03b08f445fd7ae315d3e7764347c4267f314c9ebd941a7103f0bf688c3969
+size 4522

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/session_rank_3.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c49a9ddb80d7eb729e74d7b675593b23b225bdd338b921acc6e242977979621a
+size 4544

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/session_rank_4.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4e8c48514580602d1dc6f29cb1e8e60e8c4e7ead9f9030bdaf34059b2fbf6007
+size 4533

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/session_rank_5.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:06853d983c6635b65d2200024fb3e39072e3a9b5e58fe721455be25cfffa0305
+size 4533

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/session_rank_6.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dc46b4f5ee45e9cc223bad3aad25f32557f0a932b7dd4478b8e4a35a3da32a3c
+size 4522

sess_2026_04_21_21_16_34_gcp-us2-rtx6000-blpf_petko_petkov_bridge_rich_properties_p30/session_rank_7.log

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b400a120ab22de2fe88d59d5d3a992371a8ae1ac74a1e52b630a15c592cdbe0f
+size 4522