backup: sess_2026_04_02_00_03_54_msp3-4_petko_petkov_bridge_full_tread_8b_k5 (eval outputs excluded due to 48k-file commit limit)

sess_2026_04_02_00_03_54_msp3-4_petko_petkov_bridge_full_tread_8b_k5/checkpoints/model-checkpoint-ema-030000.pt

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

sess_2026_04_02_00_03_54_msp3-4_petko_petkov_bridge_full_tread_8b_k5/hf_export/dancing-black-jellyfish/src/common_pizero_fm_paligemma.py

@@ -0,0 +1,571 @@
+from functools import cached_property
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type
+
+import torch
+import torch.distributed.fsdp
+import torch.distributed.tensor
+import torch.nn.attention.flex_attention
+import transformers
+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_02_00_03_54_msp3-4_petko_petkov_bridge_full_tread_8b_k5/hf_export/dancing-black-jellyfish/src/configuration_pizero_fm_paligemma.py

@@ -0,0 +1,366 @@
+from typing import Any, Dict, List, Optional
+
+from databib.config import Config
+
+from .common_pizero_fm_paligemma 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 PaliGemmaProcessorConfig(VLMProcessorConfig):
+    image_token: str = '<image>'
+    image_sizes: Dict[str, ImageSizeConfig] = {'main': ImageSizeConfig(width=224, height=224)}
+    max_language_tokens: int = -1
+
+    def __post_init__(self):
+        super().__post_init__()
+        for camera_name, camera_image_size in self.image_sizes.items():
+            assert camera_image_size.height % 14 == 0, f'{camera_name}: {camera_image_size}'
+            assert camera_image_size.width % 14 == 0, f'{camera_name}: {camera_image_size}'
+
+    @property
+    def num_image_tokens(self) -> Dict[str, int]:
+        return {
+            camera_name: camera_image_size.height // 14 * (camera_image_size.width // 14)
+            for (camera_name, camera_image_size) in self.image_sizes.items()
+        }
+
+    @property
+    def is_single_image_size(self) -> bool:
+        return (
+            len(self.image_sizes) == 1
+            or len(set(((image_size.height, image_size.width) for image_size in self.image_sizes.values())))
+            == 1
+        )
+
+    @property
+    def camera_names(self) -> List[str]:
+        return list(self.image_sizes.keys())
+
+
+class PaliGemmaVLMConfig(VLMConfig):
+    model_id: str = 'google/paligemma-3b-mix-224'
+    attn_implementation: str = 'flash_attention_2'
+    processor_config: PaliGemmaProcessorConfig
+    lm_head: bool = False
+    paligemma_3d_config: Dict[str, Any] = {}
+    depth_tokens: int = 0
+    train_only_depth_tokens: bool = False
+    mean_resizing: bool = False
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.train_only_depth_tokens:
+            assert self.depth_tokens > 0, self.depth_tokens
+        if self.paligemma_3d_config.get('mask_prob', 0.0) != 0.0:
+            raise NotImplementedError(
+                f"Masking is deprecated, but got mask_prob={self.paligemma_3d_config['mask_prob']}"
+            )
+
+    @property
+    def paligemma_3d_config_dict(self) -> Dict[str, Any]:
+        config = dict(self.paligemma_3d_config)
+        config['depth_config'] = dict(config['depth_config'])
+        config['depth_config']['image_sizes'] = dict(self.processor_config.image_sizes.as_json())
+        return config
+
+    @property
+    def with_depth(self) -> bool:
+        return len(self.paligemma_3d_config) > 0
+
+
+class VLAMConfig(ConfigurableModuleConfig):
+    processor_config: PiZeroFlowProcessorConfig
+    vlm_config: PaliGemmaVLMConfig
+    control_module_config: PiZeroFlowMatchingModuleConfig
+
+
+MainModelConfig = VLAMConfig

sess_2026_04_02_00_03_54_msp3-4_petko_petkov_bridge_full_tread_8b_k5/hf_export/dancing-black-jellyfish/src/format.log

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

sess_2026_04_02_00_03_54_msp3-4_petko_petkov_bridge_full_tread_8b_k5/hf_export/dancing-black-jellyfish/src/model_config.yaml

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

sess_2026_04_02_00_03_54_msp3-4_petko_petkov_bridge_full_tread_8b_k5/hf_export/dancing-black-jellyfish/src/processing_pizero_fm_paligemma.py

@@ -0,0 +1,1913 @@
+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_paligemma 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_paligemma 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_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_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_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_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_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_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