update

Author: peterdavidfagan

mujoco_robot_environments/lasa_data_generation.py

@@ -18,75 +18,112 @@
     initialize(version_base=None, config_path="./config", job_name="rearrangement")
 
     # add task configs
-    COLOR_SEPARATING_CONFIG = compose(
+    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 
+    shapes = ["CShape", "GShape", "JShape", "LShape", "NShape", "PShape", "RShape", "Sshape", "WShape", "Zshape"]
+    
+    def transform_to_workspace(trajectories, workspace_bounds):
+        """
+        Transform trajectories to fit within a workspace while maintaining aspect ratio.
+        """
+        # Extract workspace
+        workspace_x_min, workspace_x_max, workspace_y_min, workspace_y_max = workspace_bounds
 
-    def preprocess_demo(demo_data):
-        pos = demo_data.pos
-        vel = demo_data.vel 
+        # Extract range of demonstrations
+        stack = np.hstack([d.pos for d in demos])
+        original_x_min, original_y_min = np.min(stack, axis=1)
+        original_x_max, original_y_max = np.max(stack, axis=1)
+        
+        # Compute original width and height
+        original_width = original_x_max - original_x_min
+        original_height = original_y_max - original_y_min
+        
+        # Compute workspace width and height
+        workspace_width = workspace_x_max - workspace_x_min
+        workspace_height = workspace_y_max - workspace_y_min
+        
+        # Compute scaling factors
+        scale_x = workspace_width / original_width
+        scale_y = workspace_height / original_height
+        
+        # Use the smaller scaling factor to maintain aspect ratio
+        scale = min(scale_x, scale_y)
+        
+        # Compute offsets to center trajectories in the workspace
+        offset_x = workspace_x_min + (workspace_width - (original_width * scale)) / 2
+        offset_y = workspace_y_min + (workspace_height - (original_height * scale)) / 2
+        
+        # Transform each trajectory
+        transformed_position_trajectories = []
+        transformed_velocity_trajectories = []
+        for traj in trajectories:
+            pos = traj.pos
+            vel = traj.vel            
+            
+            # Scale and shift
+            transformed_x = (pos[0, :] - original_x_min) * scale + offset_x
+            transformed_y = (pos[1, :] - original_y_min) * scale + offset_y
+            z = np.repeat(0.55, pos.shape[1]) 
+            transformed_position_trajectories.append(np.vstack((transformed_x, transformed_y, z)))
 
-        # 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
+            vel_z = np.repeat(0.0, pos.shape[1]) 
+            transformed_velocity_trajectories.append(np.vstack([vel * scale / 800, vel_z])) # TODO: formalize velocity scaling
+        
+        return np.array(transformed_position_trajectories), np.array(transformed_velocity_trajectories)
 
-        # 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
+    env = LasaDrawEnv(cfg=CONFIG)
+    _, _, _, obs = env.reset()
 
-        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, joint_torques = [], [], []
-        data[f"trajectory_{demo_idx}"] = {}
+    for char in shapes:
+        demos = lasa.DataSet.__getattr__(char).demos
+        pos, vel = transform_to_workspace(demos, workspace_bounds=[0.3, 0.6, -0.3, 0.3])
+        for demo_idx, (positions, velocities) in enumerate(zip(pos, vel)):
+            print(f"Processing {char} demo {demo_idx}...")
+            joint_positions, joint_velocities, joint_torques = [], [], []
+            data[f"{char}_trajectory_{demo_idx}"] = {}
+            
+            for idx, (target_pos, target_vel) in enumerate(zip(positions.T, velocities.T)):
+                while True:
+                    pos, vel, torque = env.move_to_draw_target(target_pos, target_vel)
 
-        for idx, (target_pos, target_vel) in enumerate(zip(positions, velocities)):
-            while True:
-                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)
 
-                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:
+                        break
+                    
+                    if CONFIG.viewer:
+                        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()
 
-                # 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)
-        data[f"trajectory_{demo_idx}"]["joint_torques"] = np.vstack(joint_torques)
-    
+            # env.close()
+            data[f"{char}_trajectory_{demo_idx}"]["joint_positions"] = np.vstack(joint_positions)
+            data[f"{char}_trajectory_{demo_idx}"]["joint_velocities"] = np.vstack(joint_velocities)
+            data[f"{char}_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():

mujoco_robot_environments/lasa_policy_deployment.py

@@ -50,11 +50,15 @@ def h5_to_dict(h5_group):
     current_joint_position = starting_joint_position
 
     # load the trained model
-    with open("flax_apply_method.bin", "rb") as f:
+    # with open("ours_method.bin", "rb") as f:
+    #     serialized_from_file = f.read()
+    # model = export.deserialize(serialized_from_file)
+    # dynamics_state = jnp.zeros((1, 5000))
+    
+    with open("feedforward_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)
@@ -98,7 +102,9 @@ def preprocess_demo(demo_data):
 
     while True:
         # make a prediction using the model
-        position_target, dynamics_state = model.call(jnp.expand_dims(current_joint_position, axis=0), dynamics_state)
+        # position_target, dynamics_state = model.call(jnp.expand_dims(current_joint_position, axis=0), dynamics_state)
+        position_target = model.call(jnp.expand_dims(current_joint_position, axis=0))
+
 
         # pass the target to position actuators
         current_joint_position = env.move_to_joint_position_target(position_target[0, :7])

mujoco_robot_environments/tasks/lasa_draw.py

@@ -207,7 +207,6 @@ def reset(self, arm_configuration = None) -> dm_env.TimeStep:
         # 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, arm_configuration)
-        print(self.arm.named_configurations["home"])
 
         # configure viewer
         if self.has_viewer:

pyproject.toml

@@ -37,6 +37,7 @@ mujoco_robot_environments = ['*.png']
 [tool.poetry.dependencies]
 python = "3.10.6"
 numpy = "^1.16.0"
+jax = {extras = ["cuda12"], version = "0.4.38"}
 mujoco = "3.2.6"
 mujoco_controllers = {path="./mujoco_robot_environments/mujoco_controllers", develop=true}
 brax = "^0.10.4"