update

Author: peterdavidfagan

.gitmodules

@@ -1,3 +1,9 @@
 [submodule "mujoco_robot_environments/mujoco_controllers"]
 	path = mujoco_robot_environments/mujoco_controllers
 	url = [email protected]:peterdavidfagan/mujoco_controllers.git
+[submodule "mujoco_robot_environments/madrona_mjx"]
+	path = mujoco_robot_environments/local_builds/madrona_mjx
+	url = [email protected]:shacklettbp/madrona_mjx.git
+[submodule "mujoco_robot_environments/local_builds/robot_descriptions.py"]
+	path = mujoco_robot_environments/local_builds/robot_descriptions.py
+	url = [email protected]:peterdavidfagan/robot_descriptions.py.git

mujoco_robot_environments/config/lasa.yaml

@@ -9,6 +9,9 @@ offwidth: 640
 znear: 0.0005
 viewer: False # whether or not to use passive viewer
 
+madrona:
+  use: False
+
 # weights and biases experiment tracking
 wandb:
   use: True

mujoco_robot_environments/config/pick_mjx.yaml

@@ -0,0 +1,42 @@
+# mujoco simulation parameters
+physics_dt: 0.001
+control_dt: 0.005
+gravity: [0.0, 0.0, -9.8]
+nconmax: 1000
+njmax: 2000
+offheight: 640
+offwidth: 640
+znear: 0.0005
+viewer: False # whether or not to use passive viewer
+
+madrona:
+  use: False
+
+# weights and biases experiment tracking
+wandb:
+  use: True
+  experiment_name: rearrangement_transporter
+  tags: ["original transporter network"]
+  notes: "Curently debugging."
+  track_gradients: False
+  resume: False
+  resume_run:
+    id: xxxxx # not currently being used
+    load_epoch: 1
+  
+  metrics:
+    pick_loss:
+      step_metric: epoch
+      summary: min
+  
+  tables:
+    data_table:
+      columns: ["placeholder"]
+
+simulation_tuning_mode: False
+
+defaults:
+  - arena: lasa
+  - robots: franka_robotiq_2f85
+  - task: rearrangement
+  - dataset: default

mujoco_robot_environments/config/robots/arm/franka_emika_panda_mjx.yaml

@@ -0,0 +1,17 @@
+defaults: 
+  - actuator_config: motor
+  - sensor_config: position
+  - controller_config: osc
+
+
+default_configurations:
+  home: [0, -0.785, 0, -2.356, 0, 1.571, 0.785]
+
+arm:
+  _target_: mujoco_robot_environments.models.arms.franka_emika.FER
+  mjx: True
+  actuator_config: ${robots.arm.actuator_config}
+  sensor_config: ${robots.arm.sensor_config}
+  controller_config: ${robots.arm.controller_config}
+  configuration_config: ${robots.arm.default_configurations}
+

mujoco_robot_environments/config/robots/franka_robotiq_2f85_mjx.yaml

@@ -1,3 +1,3 @@
 defaults:
-  - arm: franka_emika_panda
+  - arm: franka_emika_panda_mjx
   - end_effector: robotiq_2f85_mjx

mujoco_robot_environments/models/arms/franka_emika.py

@@ -8,20 +8,24 @@
 from mujoco_controllers.models.arms.robot_arm import RobotArm
 
 PANDA_MJCF_PATH = os.path.join(panda_mj_description.PACKAGE_PATH, 'panda_nohand.xml')
+PANDA_MJX_MJCF_PATH = os.path.join(panda_mj_description.PACKAGE_PATH, 'mjx_panda_nohand.xml')
 
 class FER(RobotArm):
     """Franka Emika Panda Robot Arm."""
 
     def __init__(
         self, 
-        mjcf_path: str = PANDA_MJCF_PATH,
+        mjx: bool = False,
         actuator_config: dict = None, 
         sensor_config: dict = None, 
         controller_config: dict = None, 
         configuration_config: dict = None
         ):
         """Initialize the robot arm."""
-        self.mjcf_path = mjcf_path
+        if mjx:
+            self.mjcf_path = PANDA_MJX_MJCF_PATH
+        else:
+            self.mjcf_path = PANDA_MJCF_PATH
         self.actuator_config = actuator_config
         self.sensor_config = sensor_config
         self.controller_config = controller_config

mujoco_robot_environments/models/end_effectors/robotiq_2f85.py

@@ -1,13 +1,15 @@
 """Simplified Robotiq Gripper Class."""
 
+import os
 from typing import List, Tuple, Optional
 
 from dm_control import mjcf
-from robot_descriptions import robotiq_2f85_mj_description
+from robot_descriptions import robotiq_2f85_v4_mj_description
 from mujoco_controllers.models.end_effectors.robot_hand import RobotHand
 import numpy as np
 
-ROBOTIQ_MJCF_PATH = robotiq_2f85_mj_description.MJCF_PATH
+ROBOTIQ_MJCF_PATH = os.path.join(robotiq_2f85_v4_mj_description.PACKAGE_PATH, '2f85.xml')
+ROBOTIQ_MJX_MJCF_PATH = os.path.join(robotiq_2f85_v4_mj_description.PACKAGE_PATH, 'mjx_2f85.xml')
 
 class Robotiq2F85(RobotHand):
   """Robotiq 2-finger 85 adaptive gripper."""

mujoco_robot_environments/tasks/base.py

@@ -12,9 +12,7 @@
 import numpy as np
 from scipy.spatial.transform import Rotation as R
 import dm_env
-from dm_control import composer, mjcf
-from dm_control.composer.variation import distributions
-from dm_control.composer.variation import rotations
+from dm_control import mjcf
 import hydra
 from hydra.utils import instantiate
 from hydra import compose, initialize