updates to the mjcf definition (#13)

* add logic to visual geom

* add common geom setup

* fix import path

* ruff isort

sim/scripts/create_mjcf.py

@@ -5,9 +5,10 @@
     python sim/scripts/create_mjcf.py
 
 Todo:
-    0. Add IMU right position - base
-    1. Add all geoms
-    2. Condim 3 and 4 and difference in results
+    0. Add IMU to the right position
+    1. Armature damping setup for different parts of body
+    2. Test control range limits?
+    3. Add inertia in the first part of the body
 """
 
 import logging
@@ -18,7 +19,7 @@
 
 from kol.formats import mjcf
 
-from sim.env import stompy_mjcf_path
+from sim.env import model_dir, stompy_mjcf_path
 from sim.stompy.joints import StompyFixed
 
 logger = logging.getLogger(__name__)
@@ -56,6 +57,7 @@ def adapt_world(self) -> None:
                 "material", name="matplane", reflectance="0.", texture="texplane", texrepeat="1 1", texuniform="true"
             )
         )
+        asset.append(ET.Element("material", name="visualgeom", rgba="0.5 0.9 0.2 1"))
 
         compiler = root.find("compiler")
         if self.compiler is not None:
@@ -83,11 +85,6 @@ def adapt_world(self) -> None:
         # check at what stage we use this
         new_root_body.append(mjcf.Site(name="imu", size=0.01, pos=(0, 0, 0)).to_xml())
 
-        # Move gathered elements to the new root body
-        for item in items_to_move:
-            worldbody.remove(item)
-            new_root_body.append(item)
-
         # Add the new root body to the worldbody
         worldbody.append(new_root_body)
         worldbody.insert(
@@ -116,8 +113,9 @@ def adapt_world(self) -> None:
                 pos=(0.001, 0, 0),
                 quat=(1, 0, 0, 0),
                 material="matplane",
-                condim=1,
-                conaffinity=15,
+                condim=3,
+                conaffinity=1,
+                contype=0,
             ).to_xml(),
         )
 
@@ -129,7 +127,11 @@ def adapt_world(self) -> None:
             if joint in StompyFixed.default_standing().keys():
                 motors.append(
                     mjcf.Motor(
-                        name=joint, joint=joint, gear=1, ctrlrange=(limits["lower"], limits["upper"]), ctrllimited=True
+                        name=joint,
+                        joint=joint,
+                        gear=1,
+                        ctrlrange=(-200, 200),
+                        ctrllimited=True,
                     )
                 )
                 sensor_pos.append(mjcf.Actuatorpos(name=joint + "_p", actuator=joint, user="13"))
@@ -144,7 +146,6 @@ def adapt_world(self) -> None:
         sensors = root.find("sensor")
         sensors.extend(
             [
-                # TODO - pfb30 test that
                 ET.Element("framequat", name="orientation", objtype="site", noise="0.001", objname="imu"),
                 ET.Element("gyro", name="angular-velocity", site="imu", noise="0.005", cutoff="34.9"),
                 # ET.Element("framepos", name="position", objtype="site", noise="0.001", objname="imu"),
@@ -166,7 +167,10 @@ def adapt_world(self) -> None:
             ).to_xml(),
         )
 
-        # TODO - test the physical parameters
+        visual_geom = ET.Element("default", {"class": "visualgeom"})
+        geom_attributes = {"material": "visualgeom", "condim": "1", "contype": "1", "conaffinity": "0"}
+        ET.SubElement(visual_geom, "geom", geom_attributes)
+
         root.insert(
             1,
             mjcf.Default(
@@ -176,14 +180,48 @@ def adapt_world(self) -> None:
                 geom=mjcf.Geom(
                     solref=(0.001, 2),
                     friction=(0.9, 0.2, 0.2),
-                    condim=4,
+                    condim=3,
                     contype=1,
-                    conaffinity=15,
+                    conaffinity=0,
                 ),
-                # TODO - visualgem and joint param damping
+                visual_geom=visual_geom,
             ).to_xml(),
         )
-        self.tree = ET.ElementTree(root)
+
+        # Move gathered elements to the new root body
+        for item in items_to_move:
+            worldbody.remove(item)
+            new_root_body.append(item)
+
+        # add visual geom logic
+        for body in root.findall(".//body"):
+            original_geoms = list(body.findall("geom"))
+            for geom in original_geoms:
+                geom.set("class", "visualgeom")
+                # Create a new geom element
+                new_geom = ET.Element("geom")
+                new_geom.set("type", geom.get("type"))
+                new_geom.set("rgba", geom.get("rgba"))
+                new_geom.set("mesh", geom.get("mesh"))
+                if geom.get("pos"):
+                    new_geom.set("pos", geom.get("pos"))
+                if geom.get("quat"):
+                    new_geom.set("quat", geom.get("quat"))
+                new_geom.set("contype", "0")
+                new_geom.set("conaffinity", "0")
+                new_geom.set("group", "1")
+                new_geom.set("density", "0")
+
+                # Append the new geom to the body
+                index = list(body).index(geom)
+                body.insert(index + 1, new_geom)
+
+        # # Remove frc ranges?
+        # for body in root.findall(".//body"):
+        #     joints = list(body.findall("joint"))
+        #     for join in joints:
+        #         if "actuatorfrcrange" in join.attrib:
+        #             join.attrib.pop("actuatorfrcrange")
 
     def save(self, path: Union[str, Path]) -> None:
         rough_string = ET.tostring(self.tree.getroot(), "utf-8")
@@ -198,11 +236,9 @@ def save(self, path: Union[str, Path]) -> None:
     robot_name = "robot_fixed"
     robot = Sim2SimRobot(
         robot_name,
-        # TODO - fix that
-        Path("stompy"),
+        model_dir(),
         mjcf.Compiler(angle="radian", meshdir="meshes", autolimits=True),
         remove_inertia=False,
     )
-
     robot.adapt_world()
     robot.save(stompy_mjcf_path(legs_only=True))

sim/scripts/simulate_mjcf.py

@@ -37,7 +37,7 @@ def simulate(model_path: Union[str, Path], duration: float, framerate: float, re
             if time_until_next_step > 0:
                 time.sleep(time_until_next_step)
 
-            if record_video and len(frames) < data.time * framerate:
+            if record_video and (len(frames) < data.time * framerate):
                 renderer.update_scene(data)
                 pixels = renderer.render()
                 frames.append(pixels)
@@ -53,7 +53,9 @@ def simulate(model_path: Union[str, Path], duration: float, framerate: float, re
     configure_logging()
 
     parser = argparse.ArgumentParser(description="MuJoCo Simulation")
-    parser.add_argument("--model_path", type=str, default=str(stompy_mjcf_path()), help="Path to the MuJoCo XML file")
+    parser.add_argument(
+        "--model_path", type=str, default=str(stompy_mjcf_path(legs_only=True)), help="Path to the MuJoCo XML file"
+    )
     parser.add_argument("--duration", type=int, default=3, help="Duration of the simulation in seconds")
     parser.add_argument("--framerate", type=int, default=30, help="Frame rate for video recording")
     parser.add_argument("--record", action="store_true", help="Flag to record video")