Refactor UR5 environment to enhance target setting and joint control

- Updated target position and orientation initialization to use forward kinematics from the home pose.
- Enhanced `set_target` and `set_target_orientation` methods to include an option for updating joint targets via inverse kinematics.
- Modified `set_joint_positions` to optionally update the IK target based on the new end-effector pose.
- Improved the reset logic to compute the target from the end-effector's position and orientation after initialization.

robots/ur5/ur5_env.py

@@ -110,11 +110,18 @@ class UR5Env(gym.Env):
             arm_joint_names=self.JOINT_NAMES
         )
 
-        # Target position for IK - matching the new home pose
-        self._target_pos = np.array([0.4, 0.0, 0.4], dtype=np.float32)
-        # Target orientation as Euler angles (roll, pitch, yaw) in radians
-        # Default: gripper pointing down (-pi/2 pitch)
-        self._target_euler = np.array([0.0, np.pi/2, 0.0], dtype=np.float32)
+        # Set initial joint positions to home pose and compute FK
+        self.data.qpos[:6] = self.DEFAULT_HOME_POSE.copy()
+        self.data.ctrl[:6] = self.DEFAULT_HOME_POSE.copy()
+        mujoco.mj_forward(self.model, self.data)
+
+        # Target position/orientation from FK of home pose
+        ee_pos = self.get_end_effector_pos()
+        ee_quat = self.get_end_effector_quat()
+        self._target_pos = ee_pos.astype(np.float32)
+        roll, pitch, yaw = IKController.quat_to_euler(ee_quat)
+        self._target_euler = np.array([roll, pitch, yaw], dtype=np.float32)
+
         # Whether to use orientation in IK (user can toggle this)
         self._use_orientation = True
         # Gripper target: Robotiq 2F-85 uses 0=open, 255=closed
@@ -125,20 +132,49 @@ class UR5Env(gym.Env):
         # Direct joint targets (used when control_mode is 'joint')
         self._joint_targets = self.DEFAULT_HOME_POSE.copy()
 
-    def set_target(self, x: float, y: float, z: float):
-        """Set target position for IK controller."""
+    def set_target(self, x: float, y: float, z: float, update_joint_targets=True):
+        """Set target position for IK controller.
+
+        Args:
+            x, y, z: Target position in meters
+            update_joint_targets: If True, compute IK and update joint targets
+        """
         self._target_pos = np.array([x, y, z], dtype=np.float32)
         # Update mocap body position for visualization
         self.data.mocap_pos[0] = self._target_pos
 
+        # Update joint targets via IK
+        if update_joint_targets:
+            if self._use_orientation:
+                target_quat = IKController.euler_to_quat(
+                    self._target_euler[0], self._target_euler[1], self._target_euler[2]
+                )
+                self._joint_targets = self.controller.compute_ik_with_orientation(
+                    self._target_pos, target_quat
+                )
+            else:
+                self._joint_targets = self.controller.compute_ik(self._target_pos)
+
     def get_target(self):
         """Get current target position."""
         return self._target_pos.copy()
 
-    def set_target_orientation(self, roll: float, pitch: float, yaw: float):
-        """Set target orientation as Euler angles (radians)."""
+    def set_target_orientation(self, roll: float, pitch: float, yaw: float, update_joint_targets=True):
+        """Set target orientation as Euler angles (radians).
+
+        Args:
+            roll, pitch, yaw: Target orientation in radians
+            update_joint_targets: If True, compute IK and update joint targets
+        """
         self._target_euler = np.array([roll, pitch, yaw], dtype=np.float32)
 
+        # Update joint targets via IK
+        if update_joint_targets and self._use_orientation:
+            target_quat = IKController.euler_to_quat(roll, pitch, yaw)
+            self._joint_targets = self.controller.compute_ik_with_orientation(
+                self._target_pos, target_quat
+            )
+
     def get_target_orientation(self):
         """Get target orientation as Euler angles (radians)."""
         return self._target_euler.copy()
@@ -168,17 +204,36 @@ class UR5Env(gym.Env):
         """Get current control mode."""
         return self._control_mode
 
-    def set_joint_positions(self, positions):
+    def set_joint_positions(self, positions, update_ik_target=True):
         """Set target joint positions directly (bypasses IK).
 
         Args:
             positions: Array of 6 joint angles in radians
+            update_ik_target: If True, update IK target to match new EE pose
         """
         positions = np.array(positions, dtype=np.float32)
         # Clamp to joint limits
         positions = np.clip(positions, self.action_space.low[:6], self.action_space.high[:6])
         self._joint_targets = positions
 
+        # Update IK target to match the FK of new joint positions
+        if update_ik_target:
+            # Temporarily set joints to compute FK
+            old_qpos = self.data.qpos[:6].copy()
+            self.data.qpos[:6] = positions
+            mujoco.mj_forward(self.model, self.data)
+
+            # Get EE pose from FK
+            ee_pos = self.get_end_effector_pos()
+            ee_quat = self.get_end_effector_quat()
+            self._target_pos = ee_pos.astype(np.float32)
+            roll, pitch, yaw = IKController.quat_to_euler(ee_quat)
+            self._target_euler = np.array([roll, pitch, yaw], dtype=np.float32)
+            self.data.mocap_pos[0] = self._target_pos
+
+            # Restore original qpos (simulation will move towards target)
+            self.data.qpos[:6] = old_qpos
+
     def get_joint_positions(self):
         """Get current joint positions."""
         return self.data.qpos[:6].copy()
@@ -250,13 +305,21 @@ class UR5Env(gym.Env):
         self.data.qpos[box_qpos_start:box_qpos_start+3] = [0.5, 0.2, 0.45]
         self.data.qpos[box_qpos_start+3:box_qpos_start+7] = [1, 0, 0, 0]
 
-        # Reset target
-        self._target_pos = np.array([0.4, 0.0, 0.4], dtype=np.float32)
-        self._target_euler = np.array([0.0, np.pi/2, 0.0], dtype=np.float32)  # Gripper pointing down
+        # Compute forward kinematics to get EE pose from home joints
+        mujoco.mj_forward(self.model, self.data)
+
+        # Set target to current EE position (computed from FK)
+        ee_pos = self.get_end_effector_pos()
+        ee_quat = self.get_end_effector_quat()
+        self._target_pos = ee_pos.astype(np.float32)
+
+        # Convert EE quaternion to euler for target orientation
+        roll, pitch, yaw = IKController.quat_to_euler(ee_quat)
+        self._target_euler = np.array([roll, pitch, yaw], dtype=np.float32)
+
         self.data.mocap_pos[0] = self._target_pos
         self._gripper_target = 0.0  # Open gripper
-
-        mujoco.mj_forward(self.model, self.data)
+        self._joint_targets = self.DEFAULT_HOME_POSE.copy()
 
         self.steps = 0