DROID dataset transformation for actions expressed in *base* frame of the robot.
(trajectory: Dict[str, Any])
| 66 | |
| 67 | |
| 68 | def droid_baseact_transform(trajectory: Dict[str, Any]) -> Dict[str, Any]: |
| 69 | """ |
| 70 | DROID dataset transformation for actions expressed in *base* frame of the robot. |
| 71 | """ |
| 72 | dt = trajectory["action_dict"]["cartesian_velocity"][:, :3] |
| 73 | dR = trajectory["action_dict"]["cartesian_velocity"][:, 3:6] |
| 74 | |
| 75 | trajectory["action"] = tf.concat( |
| 76 | ( |
| 77 | dt, |
| 78 | dR, |
| 79 | 1 - trajectory["action_dict"]["gripper_position"], |
| 80 | ), |
| 81 | axis=-1, |
| 82 | ) |
| 83 | ( |
| 84 | trajectory["observation"]["exterior_image_1_left"], |
| 85 | trajectory["observation"]["exterior_image_2_left"], |
| 86 | ) = rand_swap_exterior_images( |
| 87 | trajectory["observation"]["exterior_image_1_left"], |
| 88 | trajectory["observation"]["exterior_image_2_left"], |
| 89 | ) |
| 90 | trajectory["observation"]["proprio"] = tf.concat( |
| 91 | ( |
| 92 | trajectory["observation"]["cartesian_position"], |
| 93 | trajectory["observation"]["gripper_position"], |
| 94 | ), |
| 95 | axis=-1, |
| 96 | ) |
| 97 | return trajectory |
| 98 | |
| 99 | |
| 100 | def droid_wristact_transform(trajectory: Dict[str, Any]) -> Dict[str, Any]: |
nothing calls this directly
no test coverage detected