855 lines
36 KiB
Python
855 lines
36 KiB
Python
"""
|
|
Adapt from https://github.com/openvla/openvla/blob/main/prismatic/vla/datasets/rlds/oxe/configs.py
|
|
configs.py
|
|
|
|
Defines per-dataset configuration (kwargs) for each dataset in Open-X Embodiment.
|
|
|
|
Configuration adopts the following structure:
|
|
image_obs_keys:
|
|
primary: primary external RGB
|
|
secondary: secondary external RGB
|
|
wrist: wrist RGB
|
|
|
|
depth_obs_keys:
|
|
primary: primary external depth
|
|
secondary: secondary external depth
|
|
wrist: wrist depth
|
|
|
|
# Always 8-dim =>> changes based on `StateEncoding`
|
|
state_obs_keys:
|
|
StateEncoding.POS_EULER: EEF XYZ (3) + Roll-Pitch-Yaw (3) + <PAD> (1) + Gripper Open/Close (1)
|
|
StateEncoding.POS_QUAT: EEF XYZ (3) + Quaternion (4) + Gripper Open/Close (1)
|
|
StateEncoding.JOINT: Joint Angles (7, <PAD> if fewer) + Gripper Open/Close (1)
|
|
|
|
state_encoding: Type of `StateEncoding`
|
|
action_encoding: Type of action encoding (e.g., EEF Position vs. Joint Position)
|
|
"""
|
|
|
|
from enum import IntEnum
|
|
from typing import Dict
|
|
|
|
import tensorflow as tf
|
|
|
|
|
|
def zero_action_filter(traj: Dict) -> bool:
|
|
"""
|
|
Filters transitions whose actions are all-0 (only relative actions, no gripper action).
|
|
Note: this filter is applied *after* action normalization, so need to compare to "normalized 0".
|
|
"""
|
|
DROID_Q01 = tf.convert_to_tensor(
|
|
[
|
|
-0.7776297926902771,
|
|
-0.5803514122962952,
|
|
-0.5795090794563293,
|
|
-0.6464047729969025,
|
|
-0.7041108310222626,
|
|
-0.8895104378461838,
|
|
]
|
|
)
|
|
DROID_Q99 = tf.convert_to_tensor(
|
|
[
|
|
0.7597932070493698,
|
|
0.5726242214441299,
|
|
0.7351000607013702,
|
|
0.6705610305070877,
|
|
0.6464948207139969,
|
|
0.8897542208433151,
|
|
]
|
|
)
|
|
DROID_NORM_0_ACT = (
|
|
2 * (tf.zeros_like(traj["action"][:, :6]) - DROID_Q01) / (DROID_Q99 - DROID_Q01 + 1e-8) - 1
|
|
)
|
|
|
|
return tf.reduce_any(tf.math.abs(traj["action"][:, :6] - DROID_NORM_0_ACT) > 1e-5)
|
|
|
|
|
|
# Defines Proprioceptive State Encoding Schemes
|
|
class StateEncoding(IntEnum):
|
|
# fmt: off
|
|
NONE = -1 # No Proprioceptive State
|
|
POS_EULER = 1 # EEF XYZ (3) + Roll-Pitch-Yaw (3) + <PAD> (1) + Gripper Open/Close (1)
|
|
POS_QUAT = 2 # EEF XYZ (3) + Quaternion (4) + Gripper Open/Close (1)
|
|
JOINT = 3 # Joint Angles (7, <PAD> if fewer) + Gripper Open/Close (1)
|
|
JOINT_BIMANUAL = 4 # Joint Angles (2 x [ Joint Angles (6) + Gripper Open/Close (1) ])
|
|
# fmt: on
|
|
|
|
|
|
# Defines Action Encoding Schemes
|
|
class ActionEncoding(IntEnum):
|
|
# fmt: off
|
|
EEF_POS = 1 # EEF Delta XYZ (3) + Roll-Pitch-Yaw (3) + Gripper Open/Close (1)
|
|
JOINT_POS = 2 # Joint Delta Position (7) + Gripper Open/Close (1)
|
|
JOINT_POS_BIMANUAL = 3 # Joint Delta Position (2 x [ Joint Delta Position (6) + Gripper Open/Close (1) ])
|
|
EEF_R6 = 4 # EEF Delta XYZ (3) + R6 (6) + Gripper Open/Close (1)
|
|
# fmt: on
|
|
|
|
|
|
# === Individual Dataset Configs ===
|
|
OXE_DATASET_CONFIGS = {
|
|
"fractal20220817_data": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["base_pose_tool_reached", "gripper_closed"],
|
|
"state_encoding": StateEncoding.POS_QUAT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 3,
|
|
"robot_type": "Google Robot",
|
|
},
|
|
"kuka": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": [
|
|
"clip_function_input/base_pose_tool_reached",
|
|
"gripper_closed",
|
|
],
|
|
"state_encoding": StateEncoding.POS_QUAT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "Kuka iiwa",
|
|
},
|
|
"bridge_oxe": { # Version of Bridge V2 in Open X-Embodiment mixture
|
|
"image_obs_keys": {"primary": "image", "secondary": "image_1", "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["EEF_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 5,
|
|
"robot_type": "WidowX",
|
|
},
|
|
"bridge_orig": { # Original version of Bridge V2 from project website
|
|
"image_obs_keys": {"primary": "image_0", "secondary": "image_1", "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["EEF_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 5,
|
|
"robot_type": "WidowX",
|
|
},
|
|
"bridge_dataset": { # Original version of Bridge V2 from project website
|
|
"image_obs_keys": {"primary": "image_0", "secondary": "image_1", "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["EEF_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 5,
|
|
"robot_type": "WidowX",
|
|
},
|
|
"taco_play": {
|
|
"image_obs_keys": {
|
|
"primary": "rgb_static",
|
|
"secondary": None,
|
|
"wrist": "rgb_gripper",
|
|
},
|
|
"depth_obs_keys": {
|
|
"primary": "depth_static",
|
|
"secondary": None,
|
|
"wrist": "depth_gripper",
|
|
},
|
|
"state_obs_keys": ["state_eef", None, "state_gripper"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 15,
|
|
"robot_type": "Franka",
|
|
},
|
|
"jaco_play": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": None,
|
|
"wrist": "image_wrist",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state_eef", None, "state_gripper"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "Jaco 2",
|
|
},
|
|
"berkeley_cable_routing": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": "top_image",
|
|
"wrist": "wrist45_image",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["robot_state", None],
|
|
"state_encoding": StateEncoding.JOINT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "Franka",
|
|
},
|
|
"roboturk": {
|
|
"image_obs_keys": {"primary": "front_rgb", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": [None, None, None, None, None, None, None, None],
|
|
"state_encoding": StateEncoding.NONE,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "Sawyer",
|
|
},
|
|
"nyu_door_opening_surprising_effectiveness": {
|
|
"image_obs_keys": {"primary": None, "secondary": None, "wrist": "image"},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": [None, None, None, None, None, None, None, None],
|
|
"state_encoding": StateEncoding.NONE,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 3,
|
|
"robot_type": "Hello Stretch",
|
|
},
|
|
"viola": {
|
|
"image_obs_keys": {
|
|
"primary": "agentview_rgb",
|
|
"secondary": None,
|
|
"wrist": "eye_in_hand_rgb",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["joint_states", "gripper_states"],
|
|
"state_encoding": StateEncoding.JOINT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 20,
|
|
"robot_type": "Franka",
|
|
},
|
|
"berkeley_autolab_ur5": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": None,
|
|
"wrist": "hand_image",
|
|
},
|
|
"depth_obs_keys": {"primary": "depth", "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state"],
|
|
"state_encoding": StateEncoding.POS_QUAT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 5,
|
|
"robot_type": "UR5",
|
|
},
|
|
"toto": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state", None],
|
|
"state_encoding": StateEncoding.JOINT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 30,
|
|
"robot_type": "Franka",
|
|
},
|
|
"language_table": {
|
|
"image_obs_keys": {"primary": "rgb", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["effector_translation", None, None, None, None, None, None],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "xArm",
|
|
},
|
|
"columbia_cairlab_pusht_real": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": None,
|
|
"wrist": "wrist_image",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["robot_state", None, None, None, None, None, None],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "UR5",
|
|
},
|
|
"stanford_kuka_multimodal_dataset_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": "depth_image", "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["ee_position", "ee_orientation", None],
|
|
"state_encoding": StateEncoding.POS_QUAT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 20,
|
|
"robot_type": "Kuka iiwa",
|
|
},
|
|
"nyu_rot_dataset_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["eef_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 3,
|
|
"robot_type": "xArm",
|
|
},
|
|
"stanford_hydra_dataset_converted_externally_to_rlds": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": None,
|
|
"wrist": "wrist_image",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["eef_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "Franka",
|
|
},
|
|
"austin_buds_dataset_converted_externally_to_rlds": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": None,
|
|
"wrist": "wrist_image",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state"],
|
|
"state_encoding": StateEncoding.JOINT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 20,
|
|
"robot_type": "Franka",
|
|
},
|
|
"nyu_franka_play_dataset_converted_externally_to_rlds": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": "image_additional_view",
|
|
"wrist": None,
|
|
},
|
|
"depth_obs_keys": {
|
|
"primary": "depth",
|
|
"secondary": "depth_additional_view",
|
|
"wrist": None,
|
|
},
|
|
"state_obs_keys": ["eef_state", None, None],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 3,
|
|
"robot_type": "Franka",
|
|
},
|
|
"maniskill_dataset_converted_externally_to_rlds": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": None,
|
|
"wrist": "wrist_image",
|
|
},
|
|
"depth_obs_keys": {
|
|
"primary": "depth",
|
|
"secondary": None,
|
|
"wrist": "wrist_depth",
|
|
},
|
|
"state_obs_keys": ["tcp_pose", "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_QUAT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 20,
|
|
"robot_type": "Franka",
|
|
},
|
|
"furniture_bench_dataset_converted_externally_to_rlds": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": None,
|
|
"wrist": "wrist_image",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state"],
|
|
"state_encoding": StateEncoding.POS_QUAT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "Franka",
|
|
},
|
|
"cmu_franka_exploration_dataset_converted_externally_to_rlds": {
|
|
"image_obs_keys": {
|
|
"primary": "highres_image",
|
|
"secondary": None,
|
|
"wrist": None,
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": [None, None, None, None, None, None, None, None],
|
|
"state_encoding": StateEncoding.NONE,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "Franka",
|
|
},
|
|
"ucsd_kitchen_dataset_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["joint_state", None],
|
|
"state_encoding": StateEncoding.JOINT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 2,
|
|
"robot_type": "xArm",
|
|
},
|
|
"ucsd_pick_and_place_dataset_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["eef_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 3,
|
|
"robot_type": "xArm",
|
|
},
|
|
"austin_sailor_dataset_converted_externally_to_rlds": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": None,
|
|
"wrist": "wrist_image",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state"],
|
|
"state_encoding": StateEncoding.POS_QUAT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 20,
|
|
"robot_type": "Franka",
|
|
},
|
|
"austin_sirius_dataset_converted_externally_to_rlds": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": None,
|
|
"wrist": "wrist_image",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state"],
|
|
"state_encoding": StateEncoding.POS_QUAT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 20,
|
|
"robot_type": "Franka",
|
|
},
|
|
"bc_z": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": [
|
|
"present/xyz",
|
|
"present/axis_angle",
|
|
None,
|
|
"present/sensed_close",
|
|
],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "Google Robot",
|
|
},
|
|
"utokyo_pr2_opening_fridge_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["eef_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "PR2",
|
|
},
|
|
"utokyo_pr2_tabletop_manipulation_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["eef_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "PR2",
|
|
},
|
|
"utokyo_xarm_pick_and_place_converted_externally_to_rlds": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": "image2",
|
|
"wrist": "hand_image",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["end_effector_pose", None, None],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "xArm",
|
|
},
|
|
"utokyo_xarm_bimanual_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["pose_r", None, None],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "xArm Bimanual",
|
|
},
|
|
"robo_net": {
|
|
"image_obs_keys": {"primary": "image", "secondary": "image1", "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["eef_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 1,
|
|
"robot_type": "Multi-Robot",
|
|
},
|
|
"berkeley_mvp_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": None, "secondary": None, "wrist": "hand_image"},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["pose", "gripper"],
|
|
"state_encoding": StateEncoding.POS_QUAT,
|
|
"action_encoding": ActionEncoding.JOINT_POS,
|
|
"control_frequency": 5,
|
|
"robot_type": "xArm",
|
|
},
|
|
"berkeley_rpt_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": None, "secondary": None, "wrist": "hand_image"},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["joint_pos", "gripper"],
|
|
"state_encoding": StateEncoding.JOINT,
|
|
"action_encoding": ActionEncoding.JOINT_POS,
|
|
"control_frequency": 30,
|
|
"robot_type": "Franka",
|
|
},
|
|
"kaist_nonprehensile_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state", None],
|
|
"state_encoding": StateEncoding.POS_QUAT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "Franka",
|
|
},
|
|
"stanford_mask_vit_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["eef_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": None,
|
|
"robot_type": "Sawyer",
|
|
},
|
|
"tokyo_u_lsmo_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["eef_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "Cobotta",
|
|
},
|
|
"dlr_sara_pour_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state", None, None],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "DLR SARA",
|
|
},
|
|
"dlr_sara_grid_clamp_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state", None, None],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "DLR SARA",
|
|
},
|
|
"dlr_edan_shared_control_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state", None],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 5,
|
|
"robot_type": "DLR EDAN",
|
|
},
|
|
"asu_table_top_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["eef_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 12.5,
|
|
"robot_type": "UR5",
|
|
},
|
|
"stanford_robocook_converted_externally_to_rlds": {
|
|
"image_obs_keys": {"primary": "image_1", "secondary": "image_2", "wrist": None},
|
|
"depth_obs_keys": {"primary": "depth_1", "secondary": "depth_2", "wrist": None},
|
|
"state_obs_keys": ["eef_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 5,
|
|
"robot_type": "Franka",
|
|
},
|
|
"imperialcollege_sawyer_wrist_cam": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": None,
|
|
"wrist": "wrist_image",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": [None, None, None, None, None, None, None, "state"],
|
|
"state_encoding": StateEncoding.NONE,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "Sawyer",
|
|
},
|
|
"iamlab_cmu_pickup_insert_converted_externally_to_rlds": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": None,
|
|
"wrist": "wrist_image",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["joint_state", "gripper_state"],
|
|
"state_encoding": StateEncoding.JOINT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 20,
|
|
"robot_type": "Franka",
|
|
},
|
|
"uiuc_d3field": {
|
|
"image_obs_keys": {"primary": "image_1", "secondary": "image_2", "wrist": None},
|
|
"depth_obs_keys": {"primary": "depth_1", "secondary": "depth_2", "wrist": None},
|
|
"state_obs_keys": [None, None, None, None, None, None, None, None],
|
|
"state_encoding": StateEncoding.NONE,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 1,
|
|
"robot_type": "Kinova Gen3",
|
|
},
|
|
"utaustin_mutex": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": None,
|
|
"wrist": "wrist_image",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state"],
|
|
"state_encoding": StateEncoding.JOINT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 20,
|
|
"robot_type": "Franka",
|
|
},
|
|
"berkeley_fanuc_manipulation": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": None,
|
|
"wrist": "wrist_image",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["joint_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.JOINT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "Fanuc Mate",
|
|
},
|
|
"cmu_playing_with_food": {
|
|
"image_obs_keys": {
|
|
"primary": "image",
|
|
"secondary": None,
|
|
"wrist": "finger_vision_1",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state", None, None],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "Franka",
|
|
},
|
|
"cmu_play_fusion": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state"],
|
|
"state_encoding": StateEncoding.JOINT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 5,
|
|
"robot_type": "Franka",
|
|
},
|
|
"cmu_stretch": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["eef_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "Hello Stretch",
|
|
},
|
|
"berkeley_gnm_recon": {
|
|
"image_obs_keys": {"primary": None, "secondary": None, "wrist": "image"},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state", None, None],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 3,
|
|
"robot_type": "Jackal",
|
|
},
|
|
"berkeley_gnm_cory_hall": {
|
|
"image_obs_keys": {"primary": None, "secondary": None, "wrist": "image"},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state", None, None],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 5,
|
|
"robot_type": "RC Car",
|
|
},
|
|
"berkeley_gnm_sac_son": {
|
|
"image_obs_keys": {"primary": None, "secondary": None, "wrist": "image"},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["state", None, None],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "TurtleBot 2",
|
|
},
|
|
# NOTE: modified
|
|
"droid": {
|
|
"image_obs_keys": {
|
|
"primary": "exterior_image_1_left",
|
|
"secondary": "exterior_image_2_left",
|
|
"wrist": "wrist_image_left",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["EEF_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_QUAT,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 15,
|
|
"robot_type": "Franka",
|
|
"aux_kwargs": {
|
|
"dataset_frame_transform_kwargs": {
|
|
"chunk_filter_fn": zero_action_filter,
|
|
},
|
|
},
|
|
},
|
|
"fmb_dataset": {
|
|
"image_obs_keys": {
|
|
"primary": "image_side_1",
|
|
"secondary": "image_side_2",
|
|
"wrist": "image_wrist_1",
|
|
},
|
|
"depth_obs_keys": {
|
|
"primary": "image_side_1_depth",
|
|
"secondary": "image_side_2_depth",
|
|
"wrist": "image_wrist_1_depth",
|
|
},
|
|
"state_obs_keys": ["proprio"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "Franka",
|
|
},
|
|
# NOTE: modified
|
|
"dobbe": {
|
|
"image_obs_keys": {"primary": "wrist_image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["EEF_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 3.75,
|
|
"robot_type": "Hello Stretch",
|
|
},
|
|
"roboset": {
|
|
"image_obs_keys": {
|
|
"primary": "image_left",
|
|
"secondary": "image_right",
|
|
"wrist": "image_wrist",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["proprio"],
|
|
"state_encoding": StateEncoding.JOINT,
|
|
"action_encoding": ActionEncoding.JOINT_POS,
|
|
"control_frequency": 5,
|
|
"robot_type": "Franka",
|
|
},
|
|
"rh20t": {
|
|
"image_obs_keys": {
|
|
"primary": "image_front",
|
|
"secondary": "image_side_right",
|
|
"wrist": "image_wrist",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["proprio"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 10,
|
|
"robot_type": "Flexiv",
|
|
},
|
|
### T-DROID datasets
|
|
"tdroid_carrot_in_bowl": { # "put carrot in bowl" task, 50 demos @ 5 Hz control
|
|
"image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["EEF_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 5,
|
|
"robot_type": "Franka",
|
|
},
|
|
"tdroid_pour_corn_in_pot": { # "pour corn from red bonawl into steel pot" task, 50 demos @ 5 Hz control
|
|
"image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["EEF_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 5,
|
|
"robot_type": "Franka",
|
|
},
|
|
"tdroid_flip_pot_upright": { # "flip pot upright" task, 10 demos @ 5 Hz control
|
|
"image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["EEF_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 5,
|
|
"robot_type": "Franka",
|
|
},
|
|
"tdroid_move_object_onto_plate": { # "move <object> onto plate" task, 150 demos @ 5 Hz control
|
|
"image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["EEF_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 5,
|
|
"robot_type": "Franka",
|
|
},
|
|
"tdroid_knock_object_over": { # "knock <object> over" task, 70 demos @ 5 Hz control
|
|
"image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["EEF_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 5,
|
|
"robot_type": "Franka",
|
|
},
|
|
"tdroid_cover_object_with_towel": { # "cover <object> with towel" task, 45 demos @ 5 Hz control
|
|
"image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None},
|
|
"depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["EEF_state", None, "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 5,
|
|
"robot_type": "Franka",
|
|
},
|
|
### DROID Finetuning datasets
|
|
"droid_wipe": {
|
|
"image_obs_keys": {
|
|
"primary": "exterior_image_2_left",
|
|
"secondary": None,
|
|
"wrist": "wrist_image_left",
|
|
},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["proprio"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 15,
|
|
"robot_type": "Franka",
|
|
},
|
|
# NOTE: modified
|
|
### LIBERO datasets (modified versions)
|
|
"libero_spatial_no_noops": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": "wrist_image"},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["EEF_state", "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 20,
|
|
"robot_type": "Franka",
|
|
},
|
|
"libero_object_no_noops": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": "wrist_image"},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["EEF_state", "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 20,
|
|
"robot_type": "Franka",
|
|
},
|
|
"libero_goal_no_noops": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": "wrist_image"},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["EEF_state", "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 20,
|
|
"robot_type": "Franka",
|
|
},
|
|
"libero_10_no_noops": {
|
|
"image_obs_keys": {"primary": "image", "secondary": None, "wrist": "wrist_image"},
|
|
"depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
|
|
"state_obs_keys": ["EEF_state", "gripper_state"],
|
|
"state_encoding": StateEncoding.POS_EULER,
|
|
"action_encoding": ActionEncoding.EEF_POS,
|
|
"control_frequency": 20,
|
|
"robot_type": "Franka",
|
|
},
|
|
}
|