Relabels actions to use reached proprioceptive state; discards last timestep (no-action).
(traj: Dict[str, Any])
| 177 | |
| 178 | # === Bridge-V2 =>> Dataset-Specific Transform === |
| 179 | def relabel_bridge_actions(traj: Dict[str, Any]) -> Dict[str, Any]: |
| 180 | """Relabels actions to use reached proprioceptive state; discards last timestep (no-action).""" |
| 181 | movement_actions = ( |
| 182 | traj["observation"]["state"][1:, :6] - traj["observation"]["state"][:-1, :6] |
| 183 | ) |
| 184 | traj_truncated = tf.nest.map_structure(lambda x: x[:-1], traj) |
| 185 | traj_truncated["action"] = tf.concat( |
| 186 | [movement_actions, traj["action"][:-1, -1:]], axis=1 |
| 187 | ) |
| 188 | |
| 189 | return traj_truncated |
| 190 | |
| 191 | |
| 192 | # === RLDS Dataset Initialization Utilities === |
no outgoing calls
no test coverage detected