update

Author: peterdavidfagan

mujoco_robot_environments/lasa_data_generation.py

@@ -55,16 +55,17 @@ def preprocess_demo(demo_data):
     data = {}
     for demo_idx, demo in enumerate(demos):
         positions, velocities = preprocess_demo(demo)
-        joint_positions, joint_velocities = [], []
+        joint_positions, joint_velocities, joint_torques = [], [], []
         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)
+                pos, vel, torque = env.move_to_draw_target(target_pos, target_vel)
 
                 if idx != 0: 
                     joint_positions.append(pos)
                     joint_velocities.append(vel)
+                    joint_torques.append(torque)
 
                 # check if target is reached
                 if env._robot.arm_controller.current_position_error() < 5e-3:
@@ -81,16 +82,18 @@ def preprocess_demo(demo_data):
                         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)
+        data[f"trajectory_{demo_idx}"]["joint_torques"] = np.vstack(joint_torques)
 
     # 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")
+            group.create_dataset("torque", data=data["joint_torques"], compression="gzip")
 
     env.close()
 

mujoco_robot_environments/lasa_policy_deployment.py

@@ -0,0 +1,123 @@
+"""Collecting robot demonstrations of LASA drawing dataset."""
+
+import jax 
+import jax.numpy as jnp
+import jax.export as export
+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=False",
+                "robots/arm/actuator_config=position",
+                ]
+                )
+
+    # load demo data to evaluate against
+    def h5_to_dict(h5_group):
+        """
+        Recursively convert an HDF5 group or dataset into a nested dictionary.
+        """
+        result = {}
+        for key, item in h5_group.items():
+            if isinstance(item, h5py.Group):  # If it's a group, recurse
+                result[key] = h5_to_dict(item)
+            elif isinstance(item, h5py.Dataset):  # If it's a dataset, load it
+                result[key] = item[:]
+        return result
+    
+    with h5py.File("./robot_trajectories.h5", "r") as f:
+        data = h5_to_dict(f)
+
+    starting_joint_position = data["trajectory_0"]["position"][0]
+    current_joint_position = starting_joint_position
+
+    # load the trained model
+    with open("flax_apply_method.bin", "rb") as f:
+        serialized_from_file = f.read()
+    model = export.deserialize(serialized_from_file)
+    dynamics_state = jnp.zeros((1, 5000))
+    
+    # instantiate the task
+    env = LasaDrawEnv(viewer=True, cfg=COLOR_SEPARATING_CONFIG) 
+    _, _, _, obs = env.reset(current_joint_position)
+
+    # draw the demonstration trajectory
+    # 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
+    
+    positions, velocities = preprocess_demo(demos[0])
+
+    for target_position in positions:
+        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_position,
+                mat=np.eye(3).flatten(),
+                rgba=[1, 0, 0, 1]
+            )
+            env.passive_view.sync()
+
+    while True:
+        # make a prediction using the model
+        position_target, dynamics_state = model.call(jnp.expand_dims(current_joint_position, axis=0), dynamics_state)
+
+        # pass the target to position actuators
+        current_joint_position = env.move_to_joint_position_target(position_target[0, :7])
+
+        # get eef xpos 
+        eef_pos = env._robot.arm_controller.current_eef_position - np.array([0.0, 0.0, 0.1])
+
+        # draw the current state
+        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=eef_pos,
+                mat=np.eye(3).flatten(),
+                rgba=[0, 1, 0, 1]
+            )
+            env.passive_view.sync()
+                    
+    env.close()
+    

mujoco_robot_environments/models/robot_arm.py

@@ -38,7 +38,8 @@ def __init__(
     self.eef_site = arm.attachment_site
     self.arm_joints = arm.joints
     self.arm_joint_ids = np.array(physics.bind(self.arm_joints).dofadr)
-    
+    self.arm_actuators = arm.actuators
+
     # set gripper and controller
     if gripper is not None:
       self.end_effector = gripper

mujoco_robot_environments/tasks/lasa_draw.py

@@ -194,16 +194,19 @@ def model(self) -> mujoco.MjModel:
     def data(self) -> mujoco.MjData:
         return self.physics.data
 
-    def reset(self) -> dm_env.TimeStep:
+    def reset(self, arm_configuration = None) -> dm_env.TimeStep:
         """Resets the environment to an initial state and returns the first
         `TimeStep` of the new episode.
         """
-        # reset the lation instance
+        if arm_configuration is None:
+            arm_configuration = self.arm.named_configurations["home"]
+
+        # reset the simulation instance
         self._physics.reset()
 
         # reset arm to home position
         # Note: for other envs we may want random sampling of initial arm positions
-        self.arm.set_joint_angles(self._physics, self.arm.named_configurations["home"])
+        self.arm.set_joint_angles(self._physics, arm_configuration)
         print(self.arm.named_configurations["home"])
 
         # configure viewer
@@ -330,7 +333,20 @@ def move_to_draw_target(self, target_position, target_velocity):
                 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
+        return self._physics.bind(self._robot.arm_joints).qpos.copy(), self._physics.bind(self._robot.arm_joints).qvel.copy(), self._physics.bind(self._robot.arm_actuators).ctrl.copy()
+
+    def move_to_joint_position_target(self, target_position):
+        """
+        Move to position and velocity target for drawing task. 
+        """
+        # step the simulation
+        for _ in range(5):
+            self._physics.set_control(target_position)
+            self._physics.step()
+            if self.passive_view is not None:
+                self.passive_view.sync()
+
+        return self._physics.bind(self._robot.arm_joints).qpos.copy()
 
 
 if __name__=="__main__":