diff --git a/serl_robot_infra/franka_env/envs/relative_env.py b/serl_robot_infra/franka_env/envs/relative_env.py index e86176ba..4b16b3b9 100644 --- a/serl_robot_infra/franka_env/envs/relative_env.py +++ b/serl_robot_infra/franka_env/envs/relative_env.py @@ -8,7 +8,7 @@ class RelativeFrame(gym.Wrapper): """ This wrapper transforms the observation and action to be expressed in the end-effector frame. Optionally, it can transform the tcp_pose into a relative frame defined as the reset pose. - + This wrapper is expected to be used on top of the base Franka environment, which has the following observation space: { @@ -56,7 +56,9 @@ def reset(self, **kwargs): self.adjoint_matrix = self.construct_adjoint_matrix(obs["state"]["tcp_pose"]) if self.include_relative_pose: # Update transformation matrix from the reset pose's relative frame to base frame - self.T_r_o_inv = np.linalg.inv(self.construct_homogeneous_matrix(obs["state"]["tcp_pose"])) + self.T_r_o_inv = np.linalg.inv( + self.construct_homogeneous_matrix(obs["state"]["tcp_pose"]) + ) # Transform observation to spatial frame return self.transform_observation(obs), info @@ -72,7 +74,7 @@ def transform_observation(self, obs): if self.include_relative_pose: T_b_o = self.construct_homogeneous_matrix(obs["state"]["tcp_pose"]) T_b_r = self.T_r_o_inv @ T_b_o - + # Reconstruct transformed tcp_pose vector p_b_r = T_b_r[:3, 3] theta_b_r = R.from_matrix(T_b_r[:3, :3]).as_quat() diff --git a/serl_robot_infra/franka_env/spacemouse/spacemouse_expert.py b/serl_robot_infra/franka_env/spacemouse/spacemouse_expert.py index 63f299ac..c7e7f412 100644 --- a/serl_robot_infra/franka_env/spacemouse/spacemouse_expert.py +++ b/serl_robot_infra/franka_env/spacemouse/spacemouse_expert.py @@ -20,7 +20,7 @@ def _read_spacemouse(self): with self.state_lock: self.latest_data["action"] = np.array( [-state.y, state.x, state.z, -state.roll, -state.pitch, -state.yaw] - ) # spacemouse axis matched with robot base frame + ) # spacemouse axis matched with robot base frame self.latest_data["buttons"] = state.buttons def get_action(self):