Function ppgm_dataset_transform

data/oxe/transforms.py:95–109

(trajectory: Dict[str, Any])

def ppgm_dataset_transform(trajectory: Dict[str, Any]) -> Dict[str, Any]:
    trajectory["action"] = tf.concat(
        [
            trajectory["action"][:, :6],
            binarize_gripper_actions(trajectory["action"][:, -1])[:, None],
        ],
        axis=1,
    )
    trajectory["observation"]["EEF_state"] = trajectory["observation"][
        "cartesian_position"
    ][:, :6]
    trajectory["observation"]["gripper_state"] = trajectory["observation"][
        "gripper_position"
    ][:, -1:]
    return trajectory


def rt1_dataset_transform(trajectory: Dict[str, Any]) -> Dict[str, Any]: