Add Florence-Bridge code and handoff documentation

977b06e verified 7 days ago

8.08 kB

"""OXE (Open-X Embodiment) benchmark — data config, embodiment tags, and mixtures."""

from starVLA.dataloader.gr00t_lerobot.datasets import ModalityConfig
from starVLA.dataloader.gr00t_lerobot.transform.base import ComposedModalityTransform
from starVLA.dataloader.gr00t_lerobot.transform.concat import ConcatTransform
from starVLA.dataloader.gr00t_lerobot.transform.state_action import (
    StateActionToTensor,
    StateActionTransform,
)
from starVLA.dataloader.gr00t_lerobot.transform.video import (
    VideoColorJitter,
    VideoCrop,
    VideoResize,
    VideoToNumpy,
    VideoToTensor,
)
from starVLA.dataloader.gr00t_lerobot.embodiment_tags import EmbodimentTag


# ---------------------------------------------------------------------------
# DataConfig — OXE Droid
# ---------------------------------------------------------------------------
class OxeDroidDataConfig:
    embodiment_tag = EmbodimentTag.OXE_DROID
    video_keys = [
        "video.exterior_image_1",
        "video.exterior_image_2",
        "video.wrist_image",
    ]
    state_keys = [
        "state.eef_position",
        "state.eef_rotation",
        "state.gripper_position",
    ]
    action_keys = [
        "action.eef_position_delta",
        "action.eef_rotation_delta",
        "action.gripper_position",
    ]
    language_keys = ["annotation.language.language_instruction"]
    observation_indices = [0]
    action_indices = list(range(16))

    def modality_config(self):
        return {
            "video": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.video_keys),
            "state": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.state_keys),
            "action": ModalityConfig(delta_indices=self.action_indices, modality_keys=self.action_keys),
            "language": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.language_keys),
        }

    def transform(self):
        return ComposedModalityTransform(transforms=[
            VideoToTensor(apply_to=self.video_keys),
            VideoCrop(apply_to=self.video_keys, scale=0.95),
            VideoResize(apply_to=self.video_keys, height=224, width=224, interpolation="linear"),
            VideoColorJitter(apply_to=self.video_keys, brightness=0.3, contrast=0.4, saturation=0.5, hue=0.08),
            VideoToNumpy(apply_to=self.video_keys),
            StateActionToTensor(apply_to=self.state_keys),
            StateActionTransform(
                apply_to=self.state_keys,
                normalization_modes={"state.eef_position": "min_max", "state.gripper_position": "min_max"},
                target_rotations={"state.eef_rotation": "rotation_6d"},
            ),
            StateActionToTensor(apply_to=self.action_keys),
            StateActionTransform(
                apply_to=self.action_keys,
                normalization_modes={"action.gripper_position": "binary"},
                target_rotations={"action.eef_rotation_delta": "axis_angle"},
            ),
            ConcatTransform(
                video_concat_order=self.video_keys,
                state_concat_order=self.state_keys,
                action_concat_order=self.action_keys,
            ),
        ])


# ---------------------------------------------------------------------------
# DataConfig — OXE Bridge
# ---------------------------------------------------------------------------
class OxeBridgeDataConfig:
    embodiment_tag = EmbodimentTag.OXE_BRIDGE
    video_keys = ["video.image_0"]
    state_keys = [
        "state.x", "state.y", "state.z",
        "state.roll", "state.pitch", "state.yaw",
        "state.pad", "state.gripper",
    ]
    action_keys = [
        "action.x", "action.y", "action.z",
        "action.roll", "action.pitch", "action.yaw",
        "action.gripper",
    ]
    language_keys = ["annotation.human.action.task_description"]
    observation_indices = [0]
    action_indices = list(range(16))

    def modality_config(self):
        return {
            "video": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.video_keys),
            "state": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.state_keys),
            "action": ModalityConfig(delta_indices=self.action_indices, modality_keys=self.action_keys),
            "language": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.language_keys),
        }

    def transform(self):
        return ComposedModalityTransform(transforms=[
            StateActionToTensor(apply_to=self.state_keys),
            StateActionTransform(
                apply_to=self.state_keys,
                normalization_modes={k: ("binary" if k == "state.gripper" else "q99") for k in self.state_keys},
            ),
            StateActionToTensor(apply_to=self.action_keys),
            StateActionTransform(
                apply_to=self.action_keys,
                normalization_modes={k: ("binary" if k == "action.gripper" else "q99") for k in self.action_keys},
            ),
        ])


# ---------------------------------------------------------------------------
# DataConfig — OXE RT-1
# ---------------------------------------------------------------------------
class OxeRT1DataConfig:
    embodiment_tag = EmbodimentTag.OXE_RT1
    video_keys = ["video.image"]
    state_keys = [
        "state.x", "state.y", "state.z",
        "state.rx", "state.ry", "state.rz", "state.rw",
        "state.gripper",
    ]
    action_keys = [
        "action.x", "action.y", "action.z",
        "action.roll", "action.pitch", "action.yaw",
        "action.gripper",
    ]
    language_keys = ["annotation.human.action.task_description"]
    observation_indices = [0]
    action_indices = list(range(16))

    def modality_config(self):
        return {
            "video": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.video_keys),
            "state": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.state_keys),
            "action": ModalityConfig(delta_indices=self.action_indices, modality_keys=self.action_keys),
            "language": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.language_keys),
        }

    def transform(self):
        return ComposedModalityTransform(transforms=[
            StateActionToTensor(apply_to=self.state_keys),
            StateActionTransform(
                apply_to=self.state_keys,
                normalization_modes={k: ("binary" if k == "state.gripper" else "q99") for k in self.state_keys},
            ),
            StateActionToTensor(apply_to=self.action_keys),
            StateActionTransform(
                apply_to=self.action_keys,
                normalization_modes={k: ("binary" if k == "action.gripper" else "q99") for k in self.action_keys},
            ),
        ])


ROBOT_TYPE_CONFIG_MAP = {
    "oxe_droid": OxeDroidDataConfig(),
    "oxe_bridge": OxeBridgeDataConfig(),
    "oxe_rt1": OxeRT1DataConfig(),
}


# ---------------------------------------------------------------------------
# Embodiment Tags
# ---------------------------------------------------------------------------
ROBOT_TYPE_TO_EMBODIMENT_TAG = {
    # Per Proposal A, embodiment_tag now lives as a classvar on each DataConfig.
    # The registry derives ROBOT_TYPE_TO_EMBODIMENT_TAG automatically. Kept as
    # an empty dict for backward compat (it is honored as legacy override).
}


# ---------------------------------------------------------------------------
# Mixtures
# ---------------------------------------------------------------------------
DATASET_NAMED_MIXTURES = {
    "bridge": [
        ("bridge_orig_1.0.0_lerobot", 1.0, "oxe_bridge"),
    ],
    "bridge_rt_1": [
        ("bridge_orig_1.0.0_lerobot", 1.0, "oxe_bridge"),
        ("fractal20220817_data_0.1.0_lerobot", 1.0, "oxe_rt1"),
    ],
    "bridge_rt_infoflow": [
        ("bridge_orig_lerobot", 1.0, "oxe_bridge"),
        ("fractal20220817_data_lerobot", 1.0, "oxe_rt1"),
    ],
}

	"""OXE (Open-X Embodiment) benchmark — data config, embodiment tags, and mixtures."""

	from starVLA.dataloader.gr00t_lerobot.datasets import ModalityConfig
	from starVLA.dataloader.gr00t_lerobot.transform.base import ComposedModalityTransform
	from starVLA.dataloader.gr00t_lerobot.transform.concat import ConcatTransform
	from starVLA.dataloader.gr00t_lerobot.transform.state_action import (
	StateActionToTensor,
	StateActionTransform,
	)
	from starVLA.dataloader.gr00t_lerobot.transform.video import (
	VideoColorJitter,
	VideoCrop,
	VideoResize,
	VideoToNumpy,
	VideoToTensor,
	)
	from starVLA.dataloader.gr00t_lerobot.embodiment_tags import EmbodimentTag


	# ---------------------------------------------------------------------------
	# DataConfig — OXE Droid
	# ---------------------------------------------------------------------------
	class OxeDroidDataConfig:
	embodiment_tag = EmbodimentTag.OXE_DROID
	video_keys = [
	"video.exterior_image_1",
	"video.exterior_image_2",
	"video.wrist_image",
	]
	state_keys = [
	"state.eef_position",
	"state.eef_rotation",
	"state.gripper_position",
	]
	action_keys = [
	"action.eef_position_delta",
	"action.eef_rotation_delta",
	"action.gripper_position",
	]
	language_keys = ["annotation.language.language_instruction"]
	observation_indices = [0]
	action_indices = list(range(16))

	def modality_config(self):
	return {
	"video": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.video_keys),
	"state": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.state_keys),
	"action": ModalityConfig(delta_indices=self.action_indices, modality_keys=self.action_keys),
	"language": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.language_keys),
	}

	def transform(self):
	return ComposedModalityTransform(transforms=[
	VideoToTensor(apply_to=self.video_keys),
	VideoCrop(apply_to=self.video_keys, scale=0.95),
	VideoResize(apply_to=self.video_keys, height=224, width=224, interpolation="linear"),
	VideoColorJitter(apply_to=self.video_keys, brightness=0.3, contrast=0.4, saturation=0.5, hue=0.08),
	VideoToNumpy(apply_to=self.video_keys),
	StateActionToTensor(apply_to=self.state_keys),
	StateActionTransform(
	apply_to=self.state_keys,
	normalization_modes={"state.eef_position": "min_max", "state.gripper_position": "min_max"},
	target_rotations={"state.eef_rotation": "rotation_6d"},
	),
	StateActionToTensor(apply_to=self.action_keys),
	StateActionTransform(
	apply_to=self.action_keys,
	normalization_modes={"action.gripper_position": "binary"},
	target_rotations={"action.eef_rotation_delta": "axis_angle"},
	),
	ConcatTransform(
	video_concat_order=self.video_keys,
	state_concat_order=self.state_keys,
	action_concat_order=self.action_keys,
	),
	])


	# ---------------------------------------------------------------------------
	# DataConfig — OXE Bridge
	# ---------------------------------------------------------------------------
	class OxeBridgeDataConfig:
	embodiment_tag = EmbodimentTag.OXE_BRIDGE
	video_keys = ["video.image_0"]
	state_keys = [
	"state.x", "state.y", "state.z",
	"state.roll", "state.pitch", "state.yaw",
	"state.pad", "state.gripper",
	]
	action_keys = [
	"action.x", "action.y", "action.z",
	"action.roll", "action.pitch", "action.yaw",
	"action.gripper",
	]
	language_keys = ["annotation.human.action.task_description"]
	observation_indices = [0]
	action_indices = list(range(16))

	def modality_config(self):
	return {
	"video": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.video_keys),
	"state": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.state_keys),
	"action": ModalityConfig(delta_indices=self.action_indices, modality_keys=self.action_keys),
	"language": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.language_keys),
	}

	def transform(self):
	return ComposedModalityTransform(transforms=[
	StateActionToTensor(apply_to=self.state_keys),
	StateActionTransform(
	apply_to=self.state_keys,
	normalization_modes={k: ("binary" if k == "state.gripper" else "q99") for k in self.state_keys},
	),
	StateActionToTensor(apply_to=self.action_keys),
	StateActionTransform(
	apply_to=self.action_keys,
	normalization_modes={k: ("binary" if k == "action.gripper" else "q99") for k in self.action_keys},
	),
	])


	# ---------------------------------------------------------------------------
	# DataConfig — OXE RT-1
	# ---------------------------------------------------------------------------
	class OxeRT1DataConfig:
	embodiment_tag = EmbodimentTag.OXE_RT1
	video_keys = ["video.image"]
	state_keys = [
	"state.x", "state.y", "state.z",
	"state.rx", "state.ry", "state.rz", "state.rw",
	"state.gripper",
	]
	action_keys = [
	"action.x", "action.y", "action.z",
	"action.roll", "action.pitch", "action.yaw",
	"action.gripper",
	]
	language_keys = ["annotation.human.action.task_description"]
	observation_indices = [0]
	action_indices = list(range(16))

	def modality_config(self):
	return {
	"video": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.video_keys),
	"state": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.state_keys),
	"action": ModalityConfig(delta_indices=self.action_indices, modality_keys=self.action_keys),
	"language": ModalityConfig(delta_indices=self.observation_indices, modality_keys=self.language_keys),
	}

	def transform(self):
	return ComposedModalityTransform(transforms=[
	StateActionToTensor(apply_to=self.state_keys),
	StateActionTransform(
	apply_to=self.state_keys,
	normalization_modes={k: ("binary" if k == "state.gripper" else "q99") for k in self.state_keys},
	),
	StateActionToTensor(apply_to=self.action_keys),
	StateActionTransform(
	apply_to=self.action_keys,
	normalization_modes={k: ("binary" if k == "action.gripper" else "q99") for k in self.action_keys},
	),
	])


	ROBOT_TYPE_CONFIG_MAP = {
	"oxe_droid": OxeDroidDataConfig(),
	"oxe_bridge": OxeBridgeDataConfig(),
	"oxe_rt1": OxeRT1DataConfig(),
	}


	# ---------------------------------------------------------------------------
	# Embodiment Tags
	# ---------------------------------------------------------------------------
	ROBOT_TYPE_TO_EMBODIMENT_TAG = {
	# Per Proposal A, embodiment_tag now lives as a classvar on each DataConfig.
	# The registry derives ROBOT_TYPE_TO_EMBODIMENT_TAG automatically. Kept as
	# an empty dict for backward compat (it is honored as legacy override).
	}


	# ---------------------------------------------------------------------------
	# Mixtures
	# ---------------------------------------------------------------------------
	DATASET_NAMED_MIXTURES = {
	"bridge": [
	("bridge_orig_1.0.0_lerobot", 1.0, "oxe_bridge"),
	],
	"bridge_rt_1": [
	("bridge_orig_1.0.0_lerobot", 1.0, "oxe_bridge"),
	("fractal20220817_data_0.1.0_lerobot", 1.0, "oxe_rt1"),
	],
	"bridge_rt_infoflow": [
	("bridge_orig_lerobot", 1.0, "oxe_bridge"),
	("fractal20220817_data_lerobot", 1.0, "oxe_rt1"),
	],
	}