update

Author: peterdavidfagan

mujoco_robot_environments/mujoco_ros_env_generation/generate_drawing_env.py

@@ -0,0 +1,27 @@
+"""A script to generate the XML file for the model."""
+
+import os
+import time
+import numpy as np
+import mujoco
+from dm_control import mjcf
+
+from mujoco_robot_environments.tasks.lasa_draw import LasaDrawEnv
+
+if __name__ == "__main__":
+    # generate the XML file for the model
+    env = LasaDrawEnv(viewer=True)
+    env.reset()
+    env._physics.step()
+    
+    # ensure cubes are in the correct position
+    mjcf_model = env._arena.mjcf_model
+    physics = mjcf.Physics.from_mjcf_model(mjcf_model)
+
+    # save as mjb
+    model_size = mujoco.mj_sizeModel(physics.model.ptr)
+    buffer = np.empty(model_size, dtype=np.uint8)
+    mujoco.mj_saveModel(physics.model.ptr, os.path.join(os.getcwd(), "lasa_draw.mjb"), None)
+
+    # sleep to allow for inspection of saved environment
+    time.sleep(10.0)

mujoco_robot_environments/tasks/lasa_draw.py

@@ -89,6 +89,7 @@ def __init__(
         self._arena.mjcf_model.visual.__getattr__("global").offwidth = cfg.offwidth
         self._arena.mjcf_model.visual.map.znear = cfg.znear
 
+
         # add table for manipulation task 
         table = Rectangle(
             name="table",
@@ -195,7 +196,7 @@ def __init__(
         self.depth_renderer = mujoco.Renderer(self._physics.model.ptr, height=self.main_camera_height, width=self.main_camera_width)
         self.depth_renderer.enable_depth_rendering()
         self.passive_view = None
-                
+                        
     def close(self) -> None:
         if self.passive_view is not None:
             self.passive_view.close()
@@ -352,9 +353,16 @@ def move_to_joint_position_target(self, target_position):
         """
         Move to position and velocity target for drawing task. 
         """
+        scale = 0
+        # sample random force and torque to apply to arm
+        force = np.random.uniform(-scale, scale, size=3) 
+        # force = np.array([scale, 0, 0])
+        torque = np.random.uniform(-scale, scale, size=3)
+
         # step the simulation
         for _ in range(5):
             self._physics.set_control(target_position)
+            self._physics.data.xfrc_applied[12] = np.concatenate([force, torque])
             self._physics.step()
             if self.passive_view is not None:
                 self.passive_view.sync()