* Fix xml path in ctrl_range and trace_hdf5_mujoco examples

* lerobotdataset load and save example scripts
* rename hdf5 file prefix
* add ee control in all environments

examples/trace_hdf5_mujoco.py → examples/hdf5_mujoco.py

@@ -11,7 +11,7 @@
 def do_replay_hdf5(args):
     # Specify the path to your HDF5 file
     with h5py.File(args.file_path, "r") as file:
-        m = mujoco.MjModel.from_xml_path("gym_lowcostrobot/assets/scene_so_arm_6dof_one_cube.xml")
+        m = mujoco.MjModel.from_xml_path("gym_lowcostrobot/assets/low_cost_robot_6dof/reach_cube.xml")
         data = mujoco.MjData(m)
         robot = SimulatedRobot(m, data)
 

examples/lerobotdataset_load.py

@@ -0,0 +1,63 @@
+import argparse, time
+import torch
+
+import mujoco
+import mujoco.viewer
+from gym_lowcostrobot.simulated_robot import SimulatedRobot
+
+import lerobot
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+
+def main():
+
+    # You can easily load a dataset from a Hugging Face repository
+    dataset = LeRobotDataset(args.repo_id)
+    print(f"\n{dataset[0]['observation.image'].shape=}")  # (4,c,h,w)
+    print(f"{dataset[0]['observation.state'].shape=}")  # (8,c)
+    print(f"{dataset[0]['action'].shape=}\n")  # (64,c)
+
+    # Finally, our datasets are fully compatible with PyTorch dataloaders 
+    # and samplers because they are just PyTorch datasets.
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        num_workers=0,
+        batch_size=32,
+        shuffle=True,
+    )
+    for batch in dataloader:
+        print(f"{batch['observation.image'].shape=}")  # (32,4,c,h,w)
+        print(f"{batch['observation.state'].shape=}")  # (32,8,c)
+        print(f"{batch['action'].shape=}")  # (32,64,c)
+        break
+
+    m = mujoco.MjModel.from_xml_path("gym_lowcostrobot/assets/scene_so_arm_6dof_one_cube.xml")
+    data = mujoco.MjData(m)
+    robot = SimulatedRobot(m, data)
+
+    with mujoco.viewer.launch_passive(m, data) as viewer:
+        # Run the simulation
+        step = 0
+        while viewer.is_running():
+            step_start = time.time()
+
+            # Step the simulation forward
+            robot.set_target_qpos(group_qpos[step][0:6])
+            mujoco.mj_step(m, data)
+
+            viewer.sync()
+
+            # Rudimentary time keeping, will drift relative to wall clock.
+            time_until_next_step = m.opt.timestep - (time.time() - step_start)
+            if time_until_next_step > 0:
+                time.sleep(time_until_next_step)
+
+            step += 1
+            print(group_qpos[step][0:6])
+            # step = step % len(group_qpos)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Replay trajectories from leRobotDataset hub")
+    parser.add_argument("--repo_id", type=str, default="thomwolf/blue_sort", help="Path to HDF5 file")
+    args = parser.parse_args()
+    main()

examples/mujoco_ctrl_range.py

@@ -2,7 +2,7 @@
 
 if __name__ == "__main__":
     # Load the model from the XML file
-    model = mujoco.MjModel.from_xml_path("gym_lowcostrobot/assets/low_cost_robot_6dof/scene_one_cube.xml")
+    model = mujoco.MjModel.from_xml_path("gym_lowcostrobot/assets/low_cost_robot_6dof/reach_cube.xml")
 
     # Extract control range for each actuator
     control_ranges = model.actuator_ctrlrange

gym_lowcostrobot/envs/lift_cube_env.py

@@ -113,6 +113,55 @@ def __init__(self, observation_mode="image", action_mode="joint", render_mode=No
         self.cube_low = np.array([-0.15, 0.10, 0.015])
         self.cube_high = np.array([0.15, 0.25, 0.015])
 
+    def inverse_kinematics(self, ee_target_pos, step=0.2, joint_name="moving_side", nb_dof=6, regularization=1e-6):
+        """
+        Computes the inverse kinematics for a robotic arm to reach the target end effector position.
+
+        :param ee_target_pos: numpy array of target end effector position [x, y, z]
+        :param step: float, step size for the iteration
+        :param joint_name: str, name of the end effector joint
+        :param nb_dof: int, number of degrees of freedom
+        :param regularization: float, regularization factor for the pseudoinverse computation
+        :return: numpy array of target joint positions
+        """
+        try:
+            # Get the joint ID from the name
+            joint_id = self.model.body(joint_name).id
+        except KeyError:
+            raise ValueError(f"Body name '{joint_name}' not found in the model.")
+
+        # Get the current end effector position
+        # ee_pos = self.d.geom_xpos[joint_id]
+        ee_id = self.model.body(joint_name).id
+        ee_pos = self.data.geom_xpos[ee_id]
+
+        # Compute the Jacobian
+        jac = np.zeros((3, self.model.nv))
+        mujoco.mj_jacBodyCom(self.model, self.data, jac, None, joint_id)
+
+        # Compute the difference between target and current end effector positions
+        delta_pos = ee_target_pos - ee_pos
+
+        # Compute the pseudoinverse of the Jacobian with regularization
+        jac_reg = jac[:, :nb_dof].T @ jac[:, :nb_dof] + regularization * np.eye(nb_dof)
+        jac_pinv = np.linalg.inv(jac_reg) @ jac[:, :nb_dof].T
+
+        # Compute target joint velocities
+        qdot = jac_pinv @ delta_pos
+
+        # Normalize joint velocities to avoid excessive movements
+        qdot_norm = np.linalg.norm(qdot)
+        if qdot_norm > 1.0:
+            qdot /= qdot_norm
+
+        # Read the current joint positions
+        qpos = self.data.qpos[:nb_dof]
+
+        # Compute the new joint positions
+        q_target_pos = qpos + qdot * step
+
+        return q_target_pos    
+
     def apply_action(self, action):
         """
         Step the simulation forward based on the action
@@ -122,7 +171,7 @@ def apply_action(self, action):
         - Joint mode: [q1, q2, q3, q4, q5, q6, gripper]
         """
         if self.action_mode == "ee":
-            raise NotImplementedError("EE mode not implemented yet")
+            #raise NotImplementedError("EE mode not implemented yet")
             ee_action, gripper_action = action[:3], action[-1]
 
             # Update the robot position based on the action

gym_lowcostrobot/envs/pick_place_cube_env.py

@@ -117,6 +117,55 @@ def __init__(self, observation_mode="image", action_mode="joint", render_mode=No
         self.target_low = np.array([-0.15, 0.10, 0.015])
         self.target_high = np.array([0.15, 0.25, 0.035])
 
+    def inverse_kinematics(self, ee_target_pos, step=0.2, joint_name="moving_side", nb_dof=6, regularization=1e-6):
+        """
+        Computes the inverse kinematics for a robotic arm to reach the target end effector position.
+
+        :param ee_target_pos: numpy array of target end effector position [x, y, z]
+        :param step: float, step size for the iteration
+        :param joint_name: str, name of the end effector joint
+        :param nb_dof: int, number of degrees of freedom
+        :param regularization: float, regularization factor for the pseudoinverse computation
+        :return: numpy array of target joint positions
+        """
+        try:
+            # Get the joint ID from the name
+            joint_id = self.model.body(joint_name).id
+        except KeyError:
+            raise ValueError(f"Body name '{joint_name}' not found in the model.")
+
+        # Get the current end effector position
+        # ee_pos = self.d.geom_xpos[joint_id]
+        ee_id = self.model.body(joint_name).id
+        ee_pos = self.data.geom_xpos[ee_id]
+
+        # Compute the Jacobian
+        jac = np.zeros((3, self.model.nv))
+        mujoco.mj_jacBodyCom(self.model, self.data, jac, None, joint_id)
+
+        # Compute the difference between target and current end effector positions
+        delta_pos = ee_target_pos - ee_pos
+
+        # Compute the pseudoinverse of the Jacobian with regularization
+        jac_reg = jac[:, :nb_dof].T @ jac[:, :nb_dof] + regularization * np.eye(nb_dof)
+        jac_pinv = np.linalg.inv(jac_reg) @ jac[:, :nb_dof].T
+
+        # Compute target joint velocities
+        qdot = jac_pinv @ delta_pos
+
+        # Normalize joint velocities to avoid excessive movements
+        qdot_norm = np.linalg.norm(qdot)
+        if qdot_norm > 1.0:
+            qdot /= qdot_norm
+
+        # Read the current joint positions
+        qpos = self.data.qpos[:nb_dof]
+
+        # Compute the new joint positions
+        q_target_pos = qpos + qdot * step
+
+        return q_target_pos
+
     def apply_action(self, action):
         """
         Step the simulation forward based on the action
@@ -126,7 +175,7 @@ def apply_action(self, action):
         - Joint mode: [q1, q2, q3, q4, q5, q6, gripper]
         """
         if self.action_mode == "ee":
-            raise NotImplementedError("EE mode not implemented yet")
+            #raise NotImplementedError("EE mode not implemented yet")
             ee_action, gripper_action = action[:3], action[-1]
 
             # Update the robot position based on the action

gym_lowcostrobot/envs/push_cube_env.py

@@ -117,6 +117,55 @@ def __init__(self, observation_mode="image", action_mode="joint", render_mode=No
         self.target_low = np.array([-0.15, 0.10, 0.015])
         self.target_high = np.array([0.15, 0.25, 0.015])
 
+    def inverse_kinematics(self, ee_target_pos, step=0.2, joint_name="moving_side", nb_dof=6, regularization=1e-6):
+        """
+        Computes the inverse kinematics for a robotic arm to reach the target end effector position.
+
+        :param ee_target_pos: numpy array of target end effector position [x, y, z]
+        :param step: float, step size for the iteration
+        :param joint_name: str, name of the end effector joint
+        :param nb_dof: int, number of degrees of freedom
+        :param regularization: float, regularization factor for the pseudoinverse computation
+        :return: numpy array of target joint positions
+        """
+        try:
+            # Get the joint ID from the name
+            joint_id = self.model.body(joint_name).id
+        except KeyError:
+            raise ValueError(f"Body name '{joint_name}' not found in the model.")
+
+        # Get the current end effector position
+        # ee_pos = self.d.geom_xpos[joint_id]
+        ee_id = self.model.body(joint_name).id
+        ee_pos = self.data.geom_xpos[ee_id]
+
+        # Compute the Jacobian
+        jac = np.zeros((3, self.model.nv))
+        mujoco.mj_jacBodyCom(self.model, self.data, jac, None, joint_id)
+
+        # Compute the difference between target and current end effector positions
+        delta_pos = ee_target_pos - ee_pos
+
+        # Compute the pseudoinverse of the Jacobian with regularization
+        jac_reg = jac[:, :nb_dof].T @ jac[:, :nb_dof] + regularization * np.eye(nb_dof)
+        jac_pinv = np.linalg.inv(jac_reg) @ jac[:, :nb_dof].T
+
+        # Compute target joint velocities
+        qdot = jac_pinv @ delta_pos
+
+        # Normalize joint velocities to avoid excessive movements
+        qdot_norm = np.linalg.norm(qdot)
+        if qdot_norm > 1.0:
+            qdot /= qdot_norm
+
+        # Read the current joint positions
+        qpos = self.data.qpos[:nb_dof]
+
+        # Compute the new joint positions
+        q_target_pos = qpos + qdot * step
+
+        return q_target_pos
+
     def apply_action(self, action):
         """
         Step the simulation forward based on the action
@@ -126,7 +175,7 @@ def apply_action(self, action):
         - Joint mode: [q1, q2, q3, q4, q5, q6, gripper]
         """
         if self.action_mode == "ee":
-            raise NotImplementedError("EE mode not implemented yet")
+            #raise NotImplementedError("EE mode not implemented yet")
             ee_action, gripper_action = action[:3], action[-1]
 
             # Update the robot position based on the action

gym_lowcostrobot/envs/reach_cube_env.py

@@ -113,6 +113,55 @@ def __init__(self, observation_mode="image", action_mode="joint", render_mode=No
         self.cube_low = np.array([-0.15, 0.10, 0.015])
         self.cube_high = np.array([0.15, 0.25, 0.015])
 
+    def inverse_kinematics(self, ee_target_pos, step=0.2, joint_name="moving_side", nb_dof=6, regularization=1e-6):
+        """
+        Computes the inverse kinematics for a robotic arm to reach the target end effector position.
+
+        :param ee_target_pos: numpy array of target end effector position [x, y, z]
+        :param step: float, step size for the iteration
+        :param joint_name: str, name of the end effector joint
+        :param nb_dof: int, number of degrees of freedom
+        :param regularization: float, regularization factor for the pseudoinverse computation
+        :return: numpy array of target joint positions
+        """
+        try:
+            # Get the joint ID from the name
+            joint_id = self.model.body(joint_name).id
+        except KeyError:
+            raise ValueError(f"Body name '{joint_name}' not found in the model.")
+
+        # Get the current end effector position
+        # ee_pos = self.d.geom_xpos[joint_id]
+        ee_id = self.model.body(joint_name).id
+        ee_pos = self.data.geom_xpos[ee_id]
+
+        # Compute the Jacobian
+        jac = np.zeros((3, self.model.nv))
+        mujoco.mj_jacBodyCom(self.model, self.data, jac, None, joint_id)
+
+        # Compute the difference between target and current end effector positions
+        delta_pos = ee_target_pos - ee_pos
+
+        # Compute the pseudoinverse of the Jacobian with regularization
+        jac_reg = jac[:, :nb_dof].T @ jac[:, :nb_dof] + regularization * np.eye(nb_dof)
+        jac_pinv = np.linalg.inv(jac_reg) @ jac[:, :nb_dof].T
+
+        # Compute target joint velocities
+        qdot = jac_pinv @ delta_pos
+
+        # Normalize joint velocities to avoid excessive movements
+        qdot_norm = np.linalg.norm(qdot)
+        if qdot_norm > 1.0:
+            qdot /= qdot_norm
+
+        # Read the current joint positions
+        qpos = self.data.qpos[:nb_dof]
+
+        # Compute the new joint positions
+        q_target_pos = qpos + qdot * step
+
+        return q_target_pos
+
     def apply_action(self, action):
         """
         Step the simulation forward based on the action
@@ -122,7 +171,7 @@ def apply_action(self, action):
         - Joint mode: [q1, q2, q3, q4, q5, q6, gripper]
         """
         if self.action_mode == "ee":
-            raise NotImplementedError("EE mode not implemented yet")
+            #raise NotImplementedError("EE mode not implemented yet")
             ee_action, gripper_action = action[:3], action[-1]
 
             # Update the robot position based on the action