add init lasa dataset

Author: peterdavidfagan

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "mujoco_robot_environments/mujoco_controllers"]
+	path = mujoco_robot_environments/mujoco_controllers
+	url = [email protected]:peterdavidfagan/mujoco_controllers.git

mujoco_robot_environments/config/robots/arm/controller_config/osc.yaml

@@ -4,8 +4,8 @@ controller_params:
   control_dt: ${control_dt}
   gains:
     position:
-      kp: 250.0
-      kd: 30.0
+      kp: 350.0
+      kd: 20.0
     orientation:
       kp: 500.0
       kd: 100.0

mujoco_robot_environments/lasa_data_generation.py

@@ -0,0 +1,96 @@
+"""Collecting robot demonstrations of LASA drawing dataset."""
+
+import numpy as np
+import mujoco
+
+import pyLasaDataset as lasa
+import h5py
+
+from mujoco_robot_environments.tasks.lasa_draw import LasaDrawEnv
+
+import hydra
+from hydra import compose, initialize
+
+
+if __name__=="__main__":
+    # clear hydra global state to avoid conflicts with other hydra instances
+    hydra.core.global_hydra.GlobalHydra.instance().clear()
+    initialize(version_base=None, config_path="./config", job_name="rearrangement")
+    
+    # add task configs
+    COLOR_SEPARATING_CONFIG = compose(
+            config_name="rearrangement",
+            overrides=[
+                "arena/props=colour_splitter",
+                "simulation_tuning_mode=True"
+                ]
+                )
+
+    # instantiate color separation task
+    env = LasaDrawEnv(viewer=True, cfg=COLOR_SEPARATING_CONFIG) 
+
+    # Leverage demonstrations from LASA dataset
+    import pyLasaDataset as lasa
+    s_data = lasa.DataSet.Sshape
+    demos = s_data.demos 
+
+    def preprocess_demo(demo_data):
+        pos = demo_data.pos
+        vel = demo_data.vel 
+
+        # scale position data
+        pos_scaled =  (pos / 200) + 0.2
+        pos_scaled = pos_scaled[:,::4]
+        positions = np.vstack([pos_scaled[1,:] + 0.2, -pos_scaled[0,:] + 0.2, np.repeat(0.55, pos_scaled.shape[1])]).T
+
+        # scale velocity data
+        vel_scaled =  (vel / 800)
+        vel_scaled = vel_scaled[:,::4]
+        velocities = np.vstack([vel_scaled[1,:], vel_scaled[0,:], np.repeat(0.0, pos_scaled.shape[1])]).T
+
+        return positions, velocities
+   
+    # interactive control of robot with mocap body
+    _, _, _, obs = env.reset()
+    data = {}
+    for demo_idx, demo in enumerate(demos):
+        positions, velocities = preprocess_demo(demo)
+        joint_positions, joint_velocities = [], []
+        data[f"trajectory_{demo_idx}"] = {}
+
+        for idx, (target_pos, target_vel) in enumerate(zip(positions, velocities)):
+            while True:
+                pos, vel = env.move_to_draw_target(target_pos, target_vel)
+
+                if idx != 0: 
+                    joint_positions.append(pos)
+                    joint_velocities.append(vel)
+
+                # check if target is reached
+                if env._robot.arm_controller.current_position_error() < 5e-3:
+                    break
+                
+                with env.passive_view.lock():
+                    env.passive_view.user_scn.ngeom += 1
+                    mujoco.mjv_initGeom(
+                        env.passive_view.user_scn.geoms[env.passive_view.user_scn.ngeom-1],
+                        type=mujoco.mjtGeom.mjGEOM_SPHERE,
+                        size=[0.001, 0, 0],
+                        pos=target_pos,
+                        mat=np.eye(3).flatten(),
+                        rgba=[1, 0, 0, 1]
+                    )
+                    env.passive_view.sync()
+                         
+        data[f"trajectory_{demo_idx}"]["joint_positions"] = np.vstack(joint_positions)
+        data[f"trajectory_{demo_idx}"]["joint_velocities"] = np.vstack(joint_velocities)
+    
+    # Save to HDF5
+    with h5py.File("robot_trajectories.h5", "w") as f:
+        for traj_name, data in data.items():
+            group = f.create_group(traj_name)
+            group.create_dataset("position", data=data["joint_positions"], compression="gzip")
+            group.create_dataset("velocity", data=data["joint_velocities"], compression="gzip")
+            
+    env.close()
+    

mujoco_robot_environments/mujoco_controllers

@@ -27,6 +27,7 @@
 from mujoco_robot_environments.environment.prop_initializer import PropPlacer
 from mujoco_robot_environments.models.robot_arm import RobotArm
 
+import h5py
 
 def generate_default_config():
     hydra.core.global_hydra.GlobalHydra.instance().clear()
@@ -40,7 +41,7 @@ def generate_default_config():
                 )
 
 
-class PushEnv(dm_env.Environment):
+class LasaDrawEnv(dm_env.Environment):
     """MuJoCo powered robotics environment with dm_env interface."""
 
     def __init__(
@@ -134,7 +135,7 @@ def __init__(
                 conaffinity=0  # no influence on collision detection
                 )
 
-        # visualise plane for drawing
+        # add a plane for drawing task
         self.draw_plane_center=self._arena.mjcf_model.worldbody.add(
                 'body',
                 name="draw_plane",
@@ -305,17 +306,17 @@ def interactive_tuning(self):
             if self.passive_view is not None:
                 self.passive_view.sync()
 
-    def test_board(self, target_position):
+    def move_to_draw_target(self, target_position, target_velocity):
         """
-        Interactively control arm to tune simulation parameters. 
+        Move to position and velocity target for drawing task. 
         """
         mocap_quat = self._physics.data.mocap_quat[0]
 
         # update control target
         self._robot.arm_controller.set_target(
-            position=target_position + [0.0, 0.0, 0.1],
+            position=target_position + [0.0, 0.0, 0.1], # account for height of toolpiece
             quat=mocap_quat, 
-            velocity=np.zeros(3),
+            velocity=target_velocity,
             angular_velocity=np.zeros(3),
             )
 
@@ -328,6 +329,9 @@ def test_board(self, target_position):
             if self.passive_view is not None:
                 self.passive_view.sync()
 
+        # return joint data for recording
+        return self._physics.bind(self._robot.arm_joints).qpos, self._physics.bind(self._robot.arm_joints).qvel
+
 
 if __name__=="__main__":
     # clear hydra global state to avoid conflicts with other hydra instances
@@ -346,26 +350,11 @@ def test_board(self, target_position):
                 )
 
     # instantiate color separation task
-    env = PushEnv(viewer=True, cfg=COLOR_SEPARATING_CONFIG) 
-    positions = [
-        np.array([0.6, 0.3, 0.55]),
-        np.array([0.6, -0.3, 0.55]),
-        np.array([0.3, -0.3, 0.55]),
-        np.array([0.3, 0.3, 0.55]),
-        np.array([0.6, -0.3, 0.55]),
-    ]
+    env = LasaDrawEnv(viewer=True, cfg=COLOR_SEPARATING_CONFIG) 
 
     # interactive control of robot with mocap body
     _, _, _, obs = env.reset()
-    for target in positions:
-        while True:
-            env.test_board(target)
-        
-            # check if target is reached
-            if env._robot.arm_controller.current_position_error() < 1e-3:
-                break
-
-        # env.interactive_tuning()
-
+    while True:
+        env.interactive_tuning()
     env.close()
 

mujoco_robot_environments/tasks/lasa_draw.py

@@ -37,10 +37,10 @@ mujoco_robot_environments = ['*.png']
 [tool.poetry.dependencies]
 python = "3.10.6"
 numpy = "^1.16.0"
-mujoco = "3.2.1"
-mujoco_controllers = "^0.0.2" 
+mujoco = "3.2.6"
+mujoco_controllers = {path="./mujoco_robot_environments/mujoco_controllers", develop=true}
 brax = "^0.10.4"
-dm-control = "1.0.22"
+dm-control = "1.0.26"
 pillow = "10.0.0"
 matplotlib = "^3.7.2"
 hydra-core = "^1.3.2"
@@ -62,6 +62,8 @@ envlogger = {extras = ["tfds"], version = "^1.2", markers = "sys_platform == 'li
 rlds = {version="^0.1.7", markers = "sys_platform == 'linux'"}
 robot-descriptions = "^1.10.0"
 mink = "^0.0.5"
+pylasadataset = "^0.1.1"
+h5py = "^3.12.1"
 
 [tool.black]
 line-length = 120