Upload folder using huggingface_hub

Browse files

Files changed (8) hide show

README.md +19 -9
assets/SO-ARM100/Simulation/SO101/scene_with_cube.xml +6 -3
assets/SO-ARM100/Simulation/SO101/so101_new_calib.xml +4 -6
assets/media/output.gif +2 -2
env.py +4 -128
src/__init__.py +0 -0
src/base_so101_env.py +129 -0
src/scripts/00_run_demo.py +45 -0

README.md CHANGED Viewed

@@ -1,22 +1,23 @@
-# Basic sO101 simulation environment.
-![output.gif](assets/media/output.gif)
-Implements a SO-101 robotic arm simulaltion environment for the [EnvHub](https://huggingface.co/docs/lerobot/envhub).
-- Observation is a `np.ndarray.shape = 640, 480`.
 - Action is a `np.ndarray.shape = 6` where each element represents the joint control.
 - Reward is the euclidian distance between the gripper and the red block, which it needs to minimize.
 ## Basic usage
 ```python
 SO101Env(
-    urdf_path=Path("assets/SO-ARM100/Simulation/SO101/scene_with_cube.xml"),
-    width=640,
-    height=480)
 env = gym.make(
     "base-sO101-env-v0",
 )
@@ -32,7 +33,15 @@ finally:
 ```
-### Assets
 All robot files are from [SO-ARM100](https://github.com/TheRobotStudio/SO-ARM100)
 ```bib
 @software{Knight_Standard_Open_SO-100,
@@ -41,6 +50,7 @@ All robot files are from [SO-ARM100](https://github.com/TheRobotStudio/SO-ARM100
     url = {https://github.com/TheRobotStudio/SO-ARM100}
 }
 ```
 [MuJoCo library used](https://github.com/google-deepmind/mujoco)
 ```bib

+# <div style="text-align:center; border-radius:30px 30px; padding:7px; color:white; margin:0; font-size:150%; font-family:Arial; background-color:#636363; overflow:hidden"><b> Basic so-101 simulation environment</b></div>
+ <p align="center">
+<img src="assets/media/output.gif" alt="Basic policy" width="75%"/>
+</p>
+Implements a so-101 robotic arm simulaltion environment for the [EnvHub](https://huggingface.co/docs/lerobot/envhub).
+- Observation is a `np.ndarray.shape = (640, 480, 3)`.
 - Action is a `np.ndarray.shape = 6` where each element represents the joint control.
 - Reward is the euclidian distance between the gripper and the red block, which it needs to minimize.
 ## Basic usage
 ```python
 SO101Env(
+    xml_pth=Path("assets/SO-ARM100/Simulation/SO101/scene_with_cube.xml"),
+    obs_w=640,
+    obs_h=480)
 env = gym.make(
     "base-sO101-env-v0",
 )
 ```
+## ToDo
+Things I want to do
+- Make a hardcoded policy that can pick up the block.
+- First attempt of training a model.
+# <div style="text-align:center; border-radius:30px 30px; padding:7px; color:white; margin:0; font-size:150%; font-family:Arial; background-color:#636363; overflow:hidden"><b> References</b></div>
+## Assets
 All robot files are from [SO-ARM100](https://github.com/TheRobotStudio/SO-ARM100)
 ```bib
 @software{Knight_Standard_Open_SO-100,
     url = {https://github.com/TheRobotStudio/SO-ARM100}
 }
 ```
+## MuJoCo
 [MuJoCo library used](https://github.com/google-deepmind/mujoco)
 ```bib

assets/SO-ARM100/Simulation/SO101/scene_with_cube.xml CHANGED Viewed

@@ -18,10 +18,13 @@
     </asset>
     <worldbody>
         <light pos="0 0 3.5" dir="0 0 -1" directional="true" />
-        <geom name="floor" size="0 0 0.05" pos="0 0 0" type="plane" material="groundplane" />
         <!-- Add a cube -->
-        <body name="cube" pos="0.25 0.01 0.01">
-            <geom name="cube_geom" type="box" size="0.02 0.01 0.01" material="cube_material" />
         </body>
     </worldbody>
 </mujoco>

     </asset>
     <worldbody>
         <light pos="0 0 3.5" dir="0 0 -1" directional="true" />
+        <geom name="floor" size="0 0 0.05" pos="0 0 0" type="plane" material="groundplane" friction="0.0 0.0 0.0"/>
         <!-- Add a cube -->
+        <body name="cube" pos="0.25 0.01 0.015">
+            <joint name="cube_joint" type="free" frictionloss="0.01"  />
+            <!-- <inertial pos="0 0 0" mass="0.025" fullinertia="6.61427e-06 1.89032e-06 5.28738e-06 -3.19807e-07 -5.90717e-09 -1.09945e-07"/> -->
+            <inertial pos="0 0 0" mass="0.0125" diaginertia="1e-5 1e-5 1e-5" />
+            <geom name="cube_geom" type="box" condim="4" solimp="2 1 0.01" solref="0.01 1" friction="5.0 4 0.01" size="0.05 0.01 0.015" material="cube_material"  mass="0.0125"/>
         </body>
     </worldbody>
 </mujoco>

assets/SO-ARM100/Simulation/SO101/so101_new_calib.xml CHANGED Viewed

@@ -20,11 +20,9 @@
     <default class="sts3215">
       <geom contype="0" conaffinity="0"/>
       <joint damping="0.60" frictionloss="0.052" armature="0.028"/>
-      <!-- For lerobot this are not exactly the motor params as the Kp and Kd not map 1-to-1 thus motor idendification with lerobot Kp=16 and Kd=32 should actually be done -->
       <position kp="17.8"/>
     </default>
     <default class="backlash">
-      <!-- +/- 0.5° of backlash -->
       <joint damping="0.01" frictionloss="0" armature="0.01" limited="true" range="-0.008726646259971648 0.008726646259971648"/>
     </default>
   </default>
@@ -102,8 +100,8 @@
                 <geom type="mesh" class="visual" pos="0.0077 0.0001 -0.0234" quat="0.707107 -0.707107 1.66015e-15 6.45094e-15" mesh="sts3215_03a_v1" material="sts3215_03a_v1_material"/>
                 <geom type="mesh" class="collision" pos="0.0077 0.0001 -0.0234" quat="0.707107 -0.707107 1.66015e-15 6.45094e-15" mesh="sts3215_03a_v1" material="sts3215_03a_v1_material"/>
                 <!-- Part wrist_roll_follower_so101_v1 -->
-                <geom type="mesh" class="visual" pos="8.32667e-17 -0.000218214 0.000949706" quat="0 1 0 0" mesh="wrist_roll_follower_so101_v1" material="wrist_roll_follower_so101_v1_material"/>
-                <geom type="mesh" class="collision" pos="8.32667e-17 -0.000218214 0.000949706" quat="0 1 0 0" mesh="wrist_roll_follower_so101_v1" material="wrist_roll_follower_so101_v1_material"/>
                 <!-- Frame gripperframe -->
                 <site group="3" name="gripperframe" pos="-0.0079 -0.000218121 -0.0981274" quat="0.707107 -0 0.707107 -2.37788e-17"/>
                 <!-- Link moving_jaw_so101_v1 -->
@@ -112,8 +110,8 @@
                   <joint axis="0 0 1" name="gripper" type="hinge" range="-0.17453297762778586 1.7453291995659765" class="sts3215"/>
                   <inertial pos="-0.00157495 -0.0300244 0.0192755" mass="0.012" fullinertia="6.61427e-06 1.89032e-06 5.28738e-06 -3.19807e-07 -5.90717e-09 -1.09945e-07"/>
                   <!-- Part moving_jaw_so101_v1 -->
-                  <geom type="mesh" class="visual" pos="-5.55112e-17 -5.55112e-17 0.0189" quat="1 -0 3.00524e-16 -2.00834e-17" mesh="moving_jaw_so101_v1" material="moving_jaw_so101_v1_material"/>
-                  <geom type="mesh" class="collision" pos="-5.55112e-17 -5.55112e-17 0.0189" quat="1 -0 3.00524e-16 -2.00834e-17" mesh="moving_jaw_so101_v1" material="moving_jaw_so101_v1_material"/>
                 </body>
               </body>
             </body>

     <default class="sts3215">
       <geom contype="0" conaffinity="0"/>
       <joint damping="0.60" frictionloss="0.052" armature="0.028"/>
       <position kp="17.8"/>
     </default>
     <default class="backlash">
       <joint damping="0.01" frictionloss="0" armature="0.01" limited="true" range="-0.008726646259971648 0.008726646259971648"/>
     </default>
   </default>
                 <geom type="mesh" class="visual" pos="0.0077 0.0001 -0.0234" quat="0.707107 -0.707107 1.66015e-15 6.45094e-15" mesh="sts3215_03a_v1" material="sts3215_03a_v1_material"/>
                 <geom type="mesh" class="collision" pos="0.0077 0.0001 -0.0234" quat="0.707107 -0.707107 1.66015e-15 6.45094e-15" mesh="sts3215_03a_v1" material="sts3215_03a_v1_material"/>
                 <!-- Part wrist_roll_follower_so101_v1 -->
+                <geom type="mesh" class="visual" pos="8.32667e-17 -0.000218214 0.000949706" quat="0 1 0 0" mesh="wrist_roll_follower_so101_v1" material="wrist_roll_follower_so101_v1_material" friction="5.0 1.0 1.0"/>
+                <geom type="mesh" class="collision" pos="8.32667e-17 -0.000218214 0.000949706" quat="0 1 0 0" mesh="wrist_roll_follower_so101_v1" material="wrist_roll_follower_so101_v1_material" friction="5.0 1.0 1.0"/>
                 <!-- Frame gripperframe -->
                 <site group="3" name="gripperframe" pos="-0.0079 -0.000218121 -0.0981274" quat="0.707107 -0 0.707107 -2.37788e-17"/>
                 <!-- Link moving_jaw_so101_v1 -->
                   <joint axis="0 0 1" name="gripper" type="hinge" range="-0.17453297762778586 1.7453291995659765" class="sts3215"/>
                   <inertial pos="-0.00157495 -0.0300244 0.0192755" mass="0.012" fullinertia="6.61427e-06 1.89032e-06 5.28738e-06 -3.19807e-07 -5.90717e-09 -1.09945e-07"/>
                   <!-- Part moving_jaw_so101_v1 -->
+                  <geom type="mesh" class="visual" pos="-5.55112e-17 -5.55112e-17 0.0189" quat="1 -0 3.00524e-16 -2.00834e-17" mesh="moving_jaw_so101_v1" material="moving_jaw_so101_v1_material" friction="5.0 0.1 0.1"/>
+                  <geom type="mesh" class="collision" pos="-5.55112e-17 -5.55112e-17 0.0189" quat="1 -0 3.00524e-16 -2.00834e-17" mesh="moving_jaw_so101_v1" material="moving_jaw_so101_v1_material" friction="5.0 0.1 0.1"/>
                 </body>
               </body>
             </body>

assets/media/output.gif CHANGED Viewed

Git LFS Details

SHA256: 7d7cc29b374f261e23336ab83960eab221764d980a2145046b5ebfb9c8e726df
Pointer size: 133 Bytes
Size of remote file: 32.5 MB

Git LFS Details

SHA256: d5e941b62c880919cb3f6adf82c40183996b81c8ec17e74f278ab4ad91987113
Pointer size: 133 Bytes
Size of remote file: 36 MB

env.py CHANGED Viewed

@@ -1,131 +1,5 @@
-from pathlib import Path
-from typing import Any
-import cv2
 import gymnasium as gym
-import mujoco
-import numpy as np
-from gymnasium import spaces
-class SO101Env(gym.Env):
-    metadata = {"render_modes": ["rgb_array"], "render_fps": 4}
-    def __init__(
-        self,
-        xml_pth: Path = Path("assets/SO-ARM100/Simulation/SO101/scene_with_cube.xml"),
-        width: int = 640,
-        height: int = 480,
-        max_episode_steps: int = 1_000,
-        camera_distance: float = 1.0,
-        camera_azimuth: int = 100,
-    ) -> None:
-        """Most simple S0101 environment. Reinforcement learning environment where reward is
-        defined by the euclidian distance between the gripper and a red block that it needs to pick up.
-        Args:
-            xml_pth (Path, optional): Path to the scene .xml file that containing the robot and the cube it needs to pickup. Defaults to Path("assets/SO-ARM100/Simulation/SO101/scene_with_cube.xml").
-            width (int, optional): Render width. Defaults to 640.
-            height (int, optional): _description_. Defaults to 480.
-            max_episode_steps (int, optional): Size of on Episode. Defaults to 200.
-            camera_distance (float, optional): Distance of the render camera to the robot. Defaults to 1.0.
-            camera_azimuth (int, optional): Azimuth of the render camera. Defaults to 100.
-        """
-        self.mj_model = mujoco.MjModel.from_xml_path(str(xml_pth))
-        self.mj_data = mujoco.MjData(self.mj_model)
-        self.width = width
-        self.height = height
-        self.mj_renderer = mujoco.Renderer(
-            self.mj_model, height=self.height, width=self.width
-        )
-        self.max_episode_steps = max_episode_steps
-        self.current_step = 0
-        self.action_space = gym.spaces.Box(
-            low=np.array([-1.91986, -1.74533, -1.69, -1.65806, -2.74385, -0.17453]),
-            high=np.array([1.91986, 1.74533, 1.69, 1.65806, 2.84121, 1.74533]),
-            shape=(self.mj_model.nu,),  # Number of actuators
-            dtype=np.float32,
-        )
-        self.observation_space = gym.spaces.Box(
-            low=0, high=255, shape=(self.height, self.width, 3), dtype=np.uint8
-        )
-        self.gripper_geom_id = mujoco.mj_name2id(
-            self.mj_model, mujoco.mjtObj.mjOBJ_BODY, "moving_jaw_so101_v1"
-        )
-        self.cube_geom_id = mujoco.mj_name2id(
-            self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, "cube_geom"
-        )
-        self.camera_distance = camera_distance
-        self.camera_azimuth = camera_azimuth
-    def step(
-        self, action: np.ndarray
-    ) -> tuple[np.ndarray, float, bool, bool, dict[str, Any]]:
-        # Apply the action and update the scene
-        self.mj_data.ctrl = action
-        mujoco.mj_step(self.mj_model, self.mj_data)
-        self.mj_renderer.update_scene(self.mj_data)
-        # Get a rendered observation
-        obs = self._get_obs()
-        # Compute the reward
-        reward = self._compute_reward()
-        # Check if the episode is terminated or truncated
-        terminated = self.current_step >= self.max_episode_steps
-        truncated = False
-        info = {}
-        self.current_step += 1
-        return obs, reward, terminated, truncated, info
-    def reset(
-        self, seed: int | None = None, options: dict[str, Any] | None = None
-    ) -> tuple[np.ndarray, dict[str, Any]]:
-        """Reset the environment to the initial state.
-        Returns:
-            np.ndarray: Initial observation
-            dict: Additional information
-        """
-        mujoco.mj_resetData(self.mj_model, self.mj_data)
-        self.current_step = 0
-        obs = self._get_obs()
-        return obs, {}
-    def _compute_reward(self) -> float:
-        """Compute the reward as the negative Euclidean distance between the gripper and the cube."""
-        # Get the positions of the gripper and cube geoms
-        gripper_pos = self.mj_data.geom_xpos[self.gripper_geom_id]
-        cube_pos = self.mj_data.geom_xpos[self.cube_geom_id]
-        # Return the negative distance as the reward
-        return -np.linalg.norm(gripper_pos - cube_pos)
-    def _get_obs(self) -> np.ndarray:
-        """Render observation
-        Returns:
-            np.ndarray: Obervation, rendered image
-        """
-        self.mj_renderer.update_scene(self.mj_data)
-        camera = mujoco.MjvCamera()
-        camera.distance = self.camera_distance  # Decrease this value to zoom in
-        camera.azimuth = self.camera_azimuth  # Camera azimuth angle (degrees)
-        self.mj_renderer.update_scene(self.mj_data, camera=camera)
-        return self.mj_renderer.render().copy()
-    def close(self) -> None:
-        # del self.mj_renderer
-        del self.mj_data
-        del self.mj_model
 gym.register(
     id="base-sO101-env-v0",
@@ -133,7 +7,9 @@ gym.register(
 )
-def make_env(n_envs: int = 1, use_async_envs: bool = False):
     """
     Create vectorized environments for your custom task.

 import gymnasium as gym
+from src.base_so101_env import SO101Env
 gym.register(
     id="base-sO101-env-v0",
 )
+def make_env(
+    n_envs: int = 1, use_async_envs: bool = False
+) -> gym.vector.AsyncVectorEnv | gym.vector.SyncVectorEnv:
     """
     Create vectorized environments for your custom task.

src/__init__.py ADDED Viewed

File without changes

src/base_so101_env.py ADDED Viewed

	@@ -0,0 +1,129 @@

+from pathlib import Path
+from typing import Any
+import cv2
+import gymnasium as gym
+import mujoco
+import numpy as np
+class SO101Env(gym.Env):
+    metadata = {"render_modes": ["rgb_array"], "render_fps": 4}
+    def __init__(
+        self,
+        xml_pth: Path = Path("assets/SO-ARM100/Simulation/SO101/scene_with_cube.xml"),
+        obs_h: int = 480,
+        obs_w: int = 640,
+        n_max_epi_steps: int = 1_000,
+        cam_dis: float = 1.0,
+        cam_azi: int = 100,
+        n_sim_steps: int = 10,
+    ) -> None:
+        """Most simple S0101 environment. Reinforcement learning environment where reward is
+        defined by the euclidian distance between the gripper and a red block that it needs to pick up.
+        Args:
+            xml_pth (Path, optional): Path to the scene .xml file that containing the robot and the cube it needs to pickup. Defaults to Path("assets/SO-ARM100/Simulation/SO101/scene_with_cube.xml").
+            obs_w (int, optional): Render obs_w. Defaults to 640.
+            obs_h (int, optional): _description_. Defaults to 480.
+            n_max_epi_steps (int, optional): Size of on Episode. Defaults to 1_000.
+            cam_dis (float, optional): Distance of the render camera to the robot. Defaults to 1.0.
+            cam_azi (int, optional): Azimuth of the render camera. Defaults to 100.
+            n_sim_steps (int, optional): Number of mujoco simulation steps.
+        """
+        self.mj_model = mujoco.MjModel.from_xml_path(str(xml_pth))
+        self.mj_data = mujoco.MjData(self.mj_model)
+        self.obs_h = obs_h
+        self.obs_w = obs_w
+        self.n_sim_steps = n_sim_steps
+        self.mj_renderer = mujoco.Renderer(
+            self.mj_model, height=self.obs_h, width=self.obs_w
+        )
+        self.n_max_epi_steps = n_max_epi_steps
+        self.current_step = 0
+        self.action_space = gym.spaces.Box(
+            low=np.array([-1.91986, -1.74533, -1.69, -1.65806, -2.74385, -0.17453]),
+            high=np.array([1.91986, 1.74533, 1.69, 1.65806, 2.84121, 1.74533]),
+            shape=(self.mj_model.nu,),
+            dtype=np.float32,
+        )
+        self.observation_space = gym.spaces.Box(
+            low=0, high=255, shape=(self.obs_h, self.obs_w, 3), dtype=np.uint8
+        )
+        self.gripper_geom_id = mujoco.mj_name2id(
+            self.mj_model, mujoco.mjtObj.mjOBJ_BODY, "moving_jaw_so101_v1"
+        )
+        self.cube_geom_id = mujoco.mj_name2id(
+            self.mj_model, mujoco.mjtObj.mjOBJ_GEOM, "cube_geom"
+        )
+        self.cam_dis = cam_dis
+        self.cam_azi = cam_azi
+    def step(
+        self, action: np.ndarray
+    ) -> tuple[np.ndarray, float, bool, bool, dict[str, Any]]:
+        # Apply the action and update the scene
+        self.mj_data.ctrl = action
+        for _ in range(self.n_sim_steps):
+            mujoco.mj_step(self.mj_model, self.mj_data)
+        self.mj_renderer.update_scene(self.mj_data)
+        # Get a rendered observation
+        obs = self._get_obs()
+        # Compute the reward
+        reward = self._compute_reward()
+        # Check if the episode is terminated or truncated
+        terminated = self.current_step >= self.n_max_epi_steps
+        truncated = False
+        info = {}
+        self.current_step += 1
+        return obs, reward, terminated, truncated, info
+    def reset(
+        self, seed: int | None = None, options: dict[str, Any] | None = None
+    ) -> tuple[np.ndarray, dict[str, Any]]:
+        """Reset the environment to the initial state.
+        Returns:
+            np.ndarray: Initial observation
+            dict: Additional information
+        """
+        mujoco.mj_resetData(self.mj_model, self.mj_data)
+        self.current_step = 0
+        obs = self._get_obs()
+        return obs, {}
+    def _compute_reward(self) -> float:
+        """Compute the reward as the negative Euclidean distance between the gripper and the cube."""
+        # Get the positions of the gripper and cube geoms
+        gripper_pos = self.mj_data.geom_xpos[self.gripper_geom_id]
+        cube_pos = self.mj_data.geom_xpos[self.cube_geom_id]
+        # Return the negative distance as the reward
+        return -np.linalg.norm(gripper_pos - cube_pos)
+    def _get_obs(self) -> np.ndarray:
+        """Render observation
+        Returns:
+            np.ndarray: Obervation, rendered image
+        """
+        mj_cam = mujoco.MjvCamera()
+        mj_cam.distance = self.cam_dis
+        mj_cam.azimuth = self.cam_azi
+        self.mj_renderer.update_scene(self.mj_data, camera=mj_cam)
+        return self.mj_renderer.render().copy()
+    def close(self) -> None:
+        # del self.mj_renderer
+        del self.mj_data
+        del self.mj_model

src/scripts/00_run_demo.py ADDED Viewed

	@@ -0,0 +1,45 @@

+from pathlib import Path
+import cv2
+import numpy as np
+import imageio
+import gymnasium as gym
+from src.base_so101_env import SO101Env
+gripper_close = 0.05
+env = SO101Env(
+    xml_pth=Path("assets/SO-ARM100/Simulation/SO101/scene_with_cube.xml"),
+    obs_w=640,
+    obs_h=480,
+    n_sim_steps=10,
+    cam_azi = 270,
+)
+frames = []   # <--- store frames here
+try:
+    obs, _ = env.reset()
+    action = np.array([0.0, 0.0, 0.0, 1, 1.5, 2])
+    # warm-up steps
+    for _ in range(10):
+        obs, reward, terminated, truncated, info = env.step(action)
+    def run_and_capture(action, steps):
+        for _ in range(steps):
+            obs, reward, terminated, truncated, info = env.step(action)
+            if terminated or truncated:
+                break
+            frames.append(obs)
+    run_and_capture(np.array([0.1, 0.2, 0.2, 1, 1.5, 2]), 10)
+    run_and_capture(np.array([0.0, 0.2, 0.2, 1, 1.5, gripper_close]), 20)
+    run_and_capture(np.array([0, -0.6, 0.2, 1, 1.5, gripper_close]), 20)
+    run_and_capture(np.array([0, -0.6, 0.1, 1, 1.5, gripper_close]), 20)
+    run_and_capture(np.array([0, -0.6, -0.0, 1, 1.5, gripper_close]), 20)
+    run_and_capture(np.array([0, -0.6, -0.2, 1, 1.5, gripper_close]), 20)
+    run_and_capture(np.array([0, -0.6, -0.4, 1, 1.5, gripper_close]), 20)
+finally:
+    env.close()
+imageio.mimsave("assets/media/output3.gif", frames, fps=20)