From bf9f36bd2dc558ef0548e290f8b1849830698803 Mon Sep 17 00:00:00 2001
From: Max <m.spahn@tudelft.nl>
Date: Wed, 24 Apr 2024 11:34:35 +0200
Subject: [PATCH] fix[mujoco]: Addresses #265. Adds tests for different mujoco
 examples.

---
 examples/.gitignore                           |  1 +
 examples/mujoco_example.py                    | 33 +++++++++++------
 examples/test_examples.py                     | 36 +++++++++++++++++--
 pyproject.toml                                |  2 +-
 urdfenvs/generic_mujoco/generic_mujoco_env.py | 10 +++++-
 5 files changed, 68 insertions(+), 14 deletions(-)

diff --git a/examples/.gitignore b/examples/.gitignore
index fb6acb5..cf5d079 100644
--- a/examples/.gitignore
+++ b/examples/.gitignore
@@ -1,3 +1,4 @@
 __pycache__
 panda
 pointRobot
+xml_model
diff --git a/examples/mujoco_example.py b/examples/mujoco_example.py
index 88c503d..34d3cfd 100644
--- a/examples/mujoco_example.py
+++ b/examples/mujoco_example.py
@@ -15,10 +15,6 @@
 from gymnasium.wrappers import RecordVideo
 
 
-ROBOTTYPE = 'panda'
-ROBOTMODEL = 'panda_without_gripper'
-ROBOTTYPE = 'pointRobot'
-ROBOTMODEL = 'pointRobot'
 
 def get_index_from_coordinates(point, mesh) -> tuple:
     distances = np.linalg.norm(mesh - point, axis=3)
@@ -40,7 +36,14 @@ def evaluate_sdf(point, mesh, sdf, resolution) -> tuple:
     return value, gradient
 
 
-def run_generic_mujoco(n_steps: int = 1000, render: Union[str, bool] = True, goal: bool = False, obstacles: bool = False):
+def run_generic_mujoco(
+        robot_name: str = 'panda',
+        robot_model: str = 'panda_without_gripper',
+        n_steps: int = 1000,
+        render: Union[str, bool] = True,
+        goal: bool = False,
+        obstacles: bool = False
+    ):
     if goal:
         goal_list = [goal1]
     else:
@@ -83,18 +86,21 @@ def run_generic_mujoco(n_steps: int = 1000, render: Union[str, bool] = True, goa
         number_constraints=1,
         physics_engine_name='mujoco',
     )
-    robot_model = RobotModel(ROBOTTYPE, ROBOTMODEL)
+    robot_model = RobotModel(robot_name, robot_model)
 
     xml_file = robot_model.get_xml_path()
     robots  = [
         GenericMujocoRobot(xml_file=xml_file, mode="vel"),
     ]
-    env: GenericMujocoEnv = gym.make(
-        'generic-mujoco-env-v0',
+    if robot_name == 'pointRobot':
+        sensors=[lidar_sensor, full_sensor, free_space_decomp, sdf_sensor]
+    else:
+        sensors=[]
+    env = GenericMujocoEnv(
         robots=robots,
         obstacles=obstacle_list,
         goals=goal_list,
-        sensors=[lidar_sensor, full_sensor, free_space_decomp, sdf_sensor],
+        sensors=sensors,
         render=render,
         enforce_real_time=True,
     ).unwrapped
@@ -120,4 +126,11 @@ def run_generic_mujoco(n_steps: int = 1000, render: Union[str, bool] = True, goa
     return history
 
 if __name__ == "__main__":
-    run_generic_mujoco(n_steps=int(1e3), render='human', obstacles=True, goal=True)
+    run_generic_mujoco(
+        robot_name='panda',
+        robot_model='panda_without_gripper',
+        n_steps=int(1e3),
+        render='human',
+        obstacles=True,
+        goal=True
+    )
diff --git a/examples/test_examples.py b/examples/test_examples.py
index ccf81f6..729039d 100644
--- a/examples/test_examples.py
+++ b/examples/test_examples.py
@@ -106,7 +106,39 @@ def test_serialization():
     from serializing_environments import run_serializing_example
     blueprint_test(run_serializing_example)
 
-def test_mujoco_example():
+def test_pointRobot_mujoco_example():
     from mujoco_example import run_generic_mujoco
-    blueprint_test(run_generic_mujoco)
+    def run_pointRobot_mujoco(
+        n_steps=int(1e3),
+        render='human',
+        obstacles=True,
+        goal=True
+    ):
+        return run_generic_mujoco(
+            robot_name='pointRobot',
+            robot_model='pointRobot',
+            n_steps=n_steps,
+            render=render,
+            obstacles=obstacles,
+            goal=goal
+        )
+    blueprint_test(run_pointRobot_mujoco)
+
+def test_panda_mujoco_example():
+    from mujoco_example import run_generic_mujoco
+    def run_panda_mujoco(
+        n_steps=int(1e3),
+        render='human',
+        obstacles=True,
+        goal=True
+    ):
+        return run_generic_mujoco(
+            robot_name='panda',
+            robot_model='panda_without_gripper',
+            n_steps=n_steps,
+            render=render,
+            obstacles=obstacles,
+            goal=goal
+        )
+    blueprint_test(run_panda_mujoco)
 
diff --git a/pyproject.toml b/pyproject.toml
index a0c0078..ac26fc7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "urdfenvs"
-version = "0.9.11"
+version = "0.9.12"
 description = "Simple simulation environment for robots, based on the urdf files."
 authors = ["Max Spahn <m.spahn@tudelft.nl>"]
 maintainers = [
diff --git a/urdfenvs/generic_mujoco/generic_mujoco_env.py b/urdfenvs/generic_mujoco/generic_mujoco_env.py
index 7284a35..4aafe67 100644
--- a/urdfenvs/generic_mujoco/generic_mujoco_env.py
+++ b/urdfenvs/generic_mujoco/generic_mujoco_env.py
@@ -165,7 +165,15 @@ def velocity_limits(self) -> np.ndarray:
     def _initialize_simulation(
         self,
     ):
-        self.model = mujoco.MjModel.from_xml_string(self._model_dm.to_xml_string())
+
+        file_name = self._xml_file.split('/')[-1]
+        mjcf.export_with_assets(
+            self._model_dm,
+            'xml_model',
+            out_file_name=file_name,
+        )
+        self.model = mujoco.MjModel.from_xml_path(f'xml_model/{file_name}')
+
 
         self.model.body_pos[0] = [0, 1, 1]
         self.model.vis.global_.offwidth = self.width