From 922477df5265d3efe046d7f8e1723ddec7c93098 Mon Sep 17 00:00:00 2001 From: Benjamin Bolte Date: Wed, 15 May 2024 18:42:22 -0700 Subject: [PATCH 1/6] static checks pass --- Makefile | 16 ++++------- pyproject.toml | 7 ++++- sim/requirements.txt | 1 + sim/scripts/create_mjcf.py | 52 +++++++++++++++++++----------------- sim/scripts/simulate_mjcf.py | 25 ++++++++++------- sim/scripts/simulate_urdf.py | 15 +++++------ sim/stompy/joints.py | 6 ++--- 7 files changed, 66 insertions(+), 56 deletions(-) diff --git a/Makefile b/Makefile index ec68abec..d82903ac 100755 --- a/Makefile +++ b/Makefile @@ -75,11 +75,9 @@ clean: # Static Checks # # ------------------------ # -py-files := $(shell find . -name '*.py') - format: - @black $(py-files) - @ruff format $(py-files) + @black sim + @ruff format sim .PHONY: format format-cpp: @@ -88,15 +86,11 @@ format-cpp: .PHONY: format-cpp static-checks: - @black --diff --check $(py-files) - @ruff check $(py-files) - @mypy --install-types --non-interactive $(py-files) + @black --diff --check sim + @ruff check sim + @mypy --install-types --non-interactive sim .PHONY: lint -mypy-daemon: - @dmypy run -- $(py-files) -.PHONY: mypy-daemon - # ------------------------ # # Unit tests # # ------------------------ # diff --git a/pyproject.toml b/pyproject.toml index d863eeb5..f5317fea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,12 +30,15 @@ warn_redundant_casts = true incremental = true namespace_packages = false +exclude = "sim/humanoid_gym/" + [[tool.mypy.overrides]] module = [ - "isaacgym.*", "humanoid.*", + "isaacgym.*", "IsaacGymEnvs.*", + "mujoco.*", ] ignore_missing_imports = true @@ -49,6 +52,8 @@ profile = "black" line-length = 120 target-version = "py310" +exclude = ["sim/humanoid_gym"] + [tool.ruff.lint] select = ["ANN", "D", "E", "F", "I", "N", "PGH", "PLC", "PLE", "PLR", "PLW", "W"] diff --git a/sim/requirements.txt b/sim/requirements.txt index 4a0b4cf3..a22a9673 100644 --- a/sim/requirements.txt +++ b/sim/requirements.txt @@ -1,3 +1,4 @@ # requirements.txt +mediapy torch diff --git a/sim/scripts/create_mjcf.py b/sim/scripts/create_mjcf.py index 90db502c..91c95bec 100644 --- a/sim/scripts/create_mjcf.py +++ b/sim/scripts/create_mjcf.py @@ -11,18 +11,21 @@ 3. Condim 3 and 4 and difference in results """ +import logging +import xml.dom.minidom import xml.etree.ElementTree as ET from pathlib import Path -import xml.dom.minidom from kol.formats import mjcf + from sim.stompy.joints import StompyFixed +logger = logging.getLogger(__name__) STOMPY_HEIGHT = 1.2 -def _pretty_print_xml(xml_string): +def _pretty_print_xml(xml_string: str) -> str: """Formats the provided XML string into a pretty-printed version.""" parsed_xml = xml.dom.minidom.parseString(xml_string) return parsed_xml.toprettyxml(indent=" ") @@ -32,26 +35,26 @@ class Sim2SimRobot(mjcf.Robot): """A class to adapt the world in a Mujoco XML file.""" def adapt_world(self) -> None: - root = self.tree.getroot() + root: ET.Element = self.tree.getroot() root.append( - mjcf.Option( + mjcf.Option( # type: ignore[call-arg] timestep=0.001, viscosity=1e-6, iterations=50, solver="PGS", gravity=(0, 0, -9.81), - flag=mjcf.Flag(frictionloss="enable"), + flag=mjcf.Flag(frictionloss="enable"), # type: ignore[attr-defined] ).to_xml() ) # TODO - test the physical parameters root.append( - mjcf.Default( - joint=mjcf.Joint(armature=0.01, damping=0.1, limited=True, frictionloss=0.01), - motor=mjcf.Motor(ctrllimited=True), - equality=mjcf.Equality(solref=(0.001, 2)), - geom=mjcf.Geom( + mjcf.Default( # type: ignore[attr-defined] + joint=mjcf.Joint(armature=0.01, damping=0.1, limited=True, frictionloss=0.01), # type: ignore[call-arg] + motor=mjcf.Motor(ctrllimited=True), # type: ignore[attr-defined] + equality=mjcf.Equality(solref=(0.001, 2)), # type: ignore[attr-defined] + geom=mjcf.Geom( # type: ignore[call-arg] solref=(0.001, 2), friction=(0.9, 0.2, 0.2), condim=4, @@ -97,16 +100,16 @@ def adapt_world(self) -> None: # Add joints to all the movement of the base new_root_body.extend( [ - mjcf.Joint(name="root_x", type="slide", axis=(1, 0, 0), limited=False).to_xml(), - mjcf.Joint(name="root_y", type="slide", axis=(0, 1, 0), limited=False).to_xml(), - mjcf.Joint(name="root_z", type="slide", axis=(0, 0, 1), limited=False).to_xml(), - mjcf.Joint(name="root_ball", type="ball", limited=False).to_xml(), + mjcf.Joint(name="root_x", type="slide", axis=(1, 0, 0), limited=False).to_xml(), # type: ignore[call-arg] + mjcf.Joint(name="root_y", type="slide", axis=(0, 1, 0), limited=False).to_xml(), # type: ignore[call-arg] + mjcf.Joint(name="root_z", type="slide", axis=(0, 0, 1), limited=False).to_xml(), # type: ignore[call-arg] + mjcf.Joint(name="root_ball", type="ball", limited=False).to_xml(), # type: ignore[call-arg] ] ) # Add imu site to the body - relative position to the body # check at what stage we use this - new_root_body.append(mjcf.Site(name="imu", size=0.01, pos=(0, 0, 0)).to_xml()) + new_root_body.append(mjcf.Site(name="imu", size=0.01, pos=(0, 0, 0)).to_xml()) # type: ignore[attr-defined] # Move gathered elements to the new root body for item in items_to_move: @@ -118,7 +121,7 @@ def adapt_world(self) -> None: worldbody.insert( 0, - mjcf.Light( + mjcf.Light( # type: ignore[attr-defined] directional=True, diffuse=(0.4, 0.4, 0.4), specular=(0.1, 0.1, 0.1), @@ -129,13 +132,13 @@ def adapt_world(self) -> None: ) worldbody.insert( 0, - mjcf.Light( + mjcf.Light( # type: ignore[attr-defined] directional=True, diffuse=(0.6, 0.6, 0.6), specular=(0.2, 0.2, 0.2), pos=(0, 0, 4), dir=(0, 0, -1) ).to_xml(), ) worldbody.insert( 0, - mjcf.Geom( + mjcf.Geom( # type: ignore[call-arg] name="ground", type="plane", size=(0, 0, 1), @@ -152,21 +155,21 @@ def adapt_world(self) -> None: for joint, limits in StompyFixed.default_limits().items(): if joint in StompyFixed.legs.all_joints(): motors.append( - mjcf.Motor( + mjcf.Motor( # type: ignore[attr-defined] name=joint, joint=joint, gear=1, ctrlrange=(limits["lower"], limits["upper"]), ctrllimited=True ) ) sensors.extend( [ - mjcf.Actuatorpos(name=joint + "_p", actuator=joint, user="13"), - mjcf.Actuatorvel(name=joint + "_v", actuator=joint, user="13"), - mjcf.Actuatorfrc(name=joint + "_f", actuator=joint, user="13", noise=0.001), + mjcf.Actuatorpos(name=joint + "_p", actuator=joint, user="13"), # type: ignore[attr-defined] + mjcf.Actuatorvel(name=joint + "_v", actuator=joint, user="13"), # type: ignore[attr-defined] + mjcf.Actuatorfrc(name=joint + "_f", actuator=joint, user="13", noise=0.001), # type: ignore[attr-defined] ] ) # Add motors and sensors - root.append(mjcf.Actuator(motors).to_xml()) - root.append(mjcf.Sensor(sensors).to_xml()) + root.append(mjcf.Actuator(motors).to_xml()) # type: ignore[arg-type, call-arg] + root.append(mjcf.Sensor(sensors).to_xml()) # type: ignore[attr-defined] # Add imus sensors = root.find("sensor") @@ -187,3 +190,4 @@ def save(self, path: str | Path) -> None: rough_string = ET.tostring(self.tree.getroot(), "utf-8") # Pretty print the XML formatted_xml = _pretty_print_xml(rough_string) + logger.info("XML:\n%s", formatted_xml) diff --git a/sim/scripts/simulate_mjcf.py b/sim/scripts/simulate_mjcf.py index 056ec51d..9efde71b 100644 --- a/sim/scripts/simulate_mjcf.py +++ b/sim/scripts/simulate_mjcf.py @@ -5,17 +5,24 @@ mjpython sim/scripts/simulate_mjcf.py --record """ +import argparse +import logging import time +from pathlib import Path + +import mediapy as media import mujoco import mujoco.viewer -import mediapy as media -import argparse +import numpy as np from sim.env import stompy_mjcf_path +from sim.logging import configure_logging + +logger = logging.getLogger(__name__) -def simulate(model_path, duration, framerate, record_video): - frames = [] +def simulate(model_path: str | Path, duration: float, framerate: float, record_video: bool) -> None: + frames: list[np.ndarray] = [] model = mujoco.MjModel.from_xml_path(model_path) data = mujoco.MjData(model) renderer = mujoco.Renderer(model) @@ -37,18 +44,18 @@ def simulate(model_path, duration, framerate, record_video): if record_video: video_path = "mjcf_simulation.mp4" media.write_video(video_path, frames, fps=framerate) - print(f"Video saved to {video_path}") + # print(f"Video saved to {video_path}") + logger.info("Video saved to %s", video_path) if __name__ == "__main__": + 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 model file" - ) + parser.add_argument("--model_path", type=str, default=str(stompy_mjcf_path()), 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") - args = parser.parse_args() simulate(args.model_path, args.duration, args.framerate, args.record) diff --git a/sim/scripts/simulate_urdf.py b/sim/scripts/simulate_urdf.py index 4b73dd4f..068e6810 100644 --- a/sim/scripts/simulate_urdf.py +++ b/sim/scripts/simulate_urdf.py @@ -8,9 +8,8 @@ """ import logging -import time from dataclasses import dataclass -from typing import Any, Dict, List, Literal, NewType +from typing import Any, Dict, Literal, NewType from isaacgym import gymapi, gymtorch, gymutil @@ -174,14 +173,14 @@ def load_gym() -> GymParams: def run_gym(gym: GymParams, mode: Literal["one_at_a_time", "all_at_once"] = "all_at_once") -> None: - joints = Stompy.all_joints() - last_time = time.time() + # joints = Stompy.all_joints() + # last_time = time.time() - dof_ids: Dict[str, int] = gym.gym.get_actor_dof_dict(gym.env, gym.robot) - body_ids: List[str] = gym.gym.get_actor_rigid_body_names(gym.env, gym.robot) + # dof_ids: Dict[str, int] = gym.gym.get_actor_dof_dict(gym.env, gym.robot) + # body_ids: List[str] = gym.gym.get_actor_rigid_body_names(gym.env, gym.robot) - joint_id = 0 - effort = 5.0 + # joint_id = 0 + # effort = 5.0 while not gym.gym.query_viewer_has_closed(gym.viewer): gym.gym.simulate(gym.sim) diff --git a/sim/stompy/joints.py b/sim/stompy/joints.py index 901bfa7f..6f4a45a6 100755 --- a/sim/stompy/joints.py +++ b/sim/stompy/joints.py @@ -7,7 +7,7 @@ import textwrap from abc import ABC -from typing import Dict, List, Union +from typing import Dict, List, Tuple, Union import numpy as np @@ -30,8 +30,8 @@ def joints(cls) -> List[str]: ] @classmethod - def joints_motors(cls) -> List[str]: - joint_names = [] + def joints_motors(cls) -> List[Tuple[str, str]]: + joint_names: List[Tuple[str, str]] = [] for attr in dir(cls): if not attr.startswith("__"): attr2 = getattr(cls, attr) From 476a5fd8ae9d1a2c3f87e1a300c80cb5c69b2b11 Mon Sep 17 00:00:00 2001 From: Benjamin Bolte Date: Wed, 15 May 2024 18:58:36 -0700 Subject: [PATCH 2/6] fix static checks --- sim/formats/__init__.py | 0 sim/formats/mjcf.py | 496 ++++++++++++++++++++++++++++++++ sim/scripts/create_mjcf.py | 58 ++-- sim/scripts/simulate_mjcf.py | 5 +- sim/scripts/update_stompy_s3.py | 10 +- 5 files changed, 534 insertions(+), 35 deletions(-) create mode 100644 sim/formats/__init__.py create mode 100644 sim/formats/mjcf.py diff --git a/sim/formats/__init__.py b/sim/formats/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/sim/formats/mjcf.py b/sim/formats/mjcf.py new file mode 100644 index 00000000..411fcf2c --- /dev/null +++ b/sim/formats/mjcf.py @@ -0,0 +1,496 @@ +# mypy: disable-error-code="operator,union-attr" +"""Defines common types and functions for exporting MJCF files. + +API reference: +https://github.com/google-deepmind/mujoco/blob/main/src/xml/xml_native_writer.cc#L780 + +Todo: + 1. Clean up the inertia config + 2. Add visual and imu support +""" + +import glob +import os +import shutil +import xml.etree.ElementTree as ET +from dataclasses import dataclass, field +from pathlib import Path +from typing import List, Literal, Optional, Tuple, Union + +import mujoco + + +@dataclass +class Compiler: + coordinate: Optional[Literal["local", "global"]] = None + angle: Literal["radian", "degree"] = "radian" + meshdir: Optional[str] = None + eulerseq: Optional[Literal["xyz", "zxy", "zyx", "yxz", "yzx", "xzy"]] = None + autolimits: Optional[bool] = None + + def to_xml(self, compiler: Optional[ET.Element] = None) -> ET.Element: + if compiler is None: + compiler = ET.Element("compiler") + if self.coordinate is not None: + compiler.set("coordinate", self.coordinate) + compiler.set("angle", self.angle) + if self.meshdir is not None: + compiler.set("meshdir", self.meshdir) + if self.eulerseq is not None: + compiler.set("eulerseq", self.eulerseq) + if self.autolimits is not None: + compiler.set("autolimits", str(self.autolimits).lower()) + return compiler + + +@dataclass +class Mesh: + name: str + file: str + scale: Optional[Tuple[float, float, float]] = None + + def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: + mesh = ET.Element("mesh") if root is None else ET.SubElement(root, "mesh") + mesh.set("name", self.name) + mesh.set("file", self.file) + if self.scale is not None: + mesh.set("scale", " ".join(map(str, self.scale))) + return mesh + + +@dataclass +class Joint: + name: Optional[str] = None + type: Optional[Literal["hinge", "slide", "ball", "free"]] = None + pos: Optional[Tuple[float, float, float]] = None + axis: Optional[Tuple[float, float, float]] = None + limited: Optional[bool] = None + range: Optional[Tuple[float, float]] = None + damping: Optional[float] = None + stiffness: Optional[float] = None + armature: Optional[float] = None + frictionloss: Optional[float] = None + + def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: + joint = ET.Element("joint") if root is None else ET.SubElement(root, "joint") + if self.name is not None: + joint.set("name", self.name) + if self.type is not None: + joint.set("type", self.type) + if self.pos is not None: + joint.set("pos", " ".join(map(str, self.pos))) + if self.axis is not None: + joint.set("axis", " ".join(map(str, self.axis))) + if self.range is not None: + joint.set("range", " ".join(map(str, self.range))) + if self.limited is not None: + joint.set("limited", str(self.limited).lower()) + if self.damping is not None: + joint.set("damping", str(self.damping)) + if self.stiffness is not None: + joint.set("stiffness", str(self.stiffness)) + if self.armature is not None: + joint.set("armature", str(self.armature)) + if self.frictionloss is not None: + joint.set("frictionloss", str(self.frictionloss)) + return joint + + +@dataclass +class Geom: + name: Optional[str] = None + type: Optional[Literal["plane", "sphere", "cylinder", "box", "capsule", "ellipsoid", "mesh"]] = None + plane: Optional[str] = None + rgba: Optional[Tuple[float, float, float, float]] = None + pos: Optional[Tuple[float, float, float]] = None + quat: Optional[Tuple[float, float, float, float]] = None + matplane: Optional[str] = None + material: Optional[str] = None + condim: Optional[int] = None + contype: Optional[int] = None + conaffinity: Optional[int] = None + size: Optional[Tuple[float, float, float]] = None + friction: Optional[Tuple[float, float, float]] = None + solref: Optional[Tuple[float, float]] = None + density: Optional[float] = None + + def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: + geom = ET.Element("geom") if root is None else ET.SubElement(root, "geom") + if self.name is not None: + geom.set("name", self.name) + if self.type is not None: + geom.set("type", self.type) + if self.rgba is not None: + geom.set("rgba", " ".join(map(str, self.rgba))) + if self.pos is not None: + geom.set("pos", " ".join(map(str, self.pos))) + if self.quat is not None: + geom.set("quat", " ".join(map(str, self.quat))) + if self.matplane is not None: + geom.set("matplane", self.matplane) + if self.material is not None: + geom.set("material", self.material) + if self.condim is not None: + geom.set("condim", str(self.condim)) + if self.contype is not None: + geom.set("contype", str(self.contype)) + if self.conaffinity is not None: + geom.set("conaffinity", str(self.conaffinity)) + if self.plane is not None: + geom.set("plane", self.plane) + if self.size is not None: + geom.set("size", " ".join(map(str, self.size))) + if self.friction is not None: + geom.set("friction", " ".join(map(str, self.friction))) + if self.solref is not None: + geom.set("solref", " ".join(map(str, self.solref))) + if self.density is not None: + geom.set("density", str(self.density)) + return geom + + +@dataclass +class Body: + name: str + pos: Optional[Tuple[float, float, float]] = field(default=None) + quat: Optional[Tuple[float, float, float, float]] = field(default=None) + geom: Optional[Geom] = field(default=None) + joint: Optional[Joint] = field(default=None) + # TODO - fix inertia, until then rely on Mujoco's engine + # inertial: Inertial = None + + def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: + body = ET.Element("body") if root is None else ET.SubElement(root, "body") + body.set("name", self.name) + if self.pos is not None: + body.set("pos", " ".join(map(str, self.pos))) + if self.quat is not None: + body.set("quat", " ".join(f"{x:.8g}" for x in self.quat)) + if self.joint is not None: + self.joint.to_xml(body) + if self.geom is not None: + self.geom.to_xml(body) + return body + + +@dataclass +class Flag: + frictionloss: Optional[str] = None + # removed at 3.1.4 + # sensornoise: Optional[str] = None + + def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: + flag = ET.Element("flag") if root is None else ET.SubElement(root, "flag") + if self.frictionloss is not None: + flag.set("frictionloss", self.frictionloss) + return flag + + +@dataclass +class Option: + timestep: Optional[float] = None + viscosity: Optional[float] = None + iterations: Optional[int] = None + solver: Optional[Literal["PGS", "CG", "Newton"]] = None + gravity: Optional[Tuple[float, float, float]] = None + flag: Optional[Flag] = None + + def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: + if root is None: + option = ET.Element("option") + else: + option = ET.SubElement(root, "option") + if self.iterations is not None: + option.set("iterations", str(self.iterations)) + if self.timestep is not None: + option.set("timestep", str(self.timestep)) + if self.viscosity is not None: + option.set("viscosity", str(self.viscosity)) + if self.solver is not None: + option.set("solver", self.solver) + if self.gravity is not None: + option.set("gravity", " ".join(map(str, self.gravity))) + if self.flag is not None: + self.flag.to_xml(option) + return option + + +@dataclass +class Motor: + name: Optional[str] = None + joint: Optional[str] = None + ctrlrange: Optional[Tuple[float, float]] = None + ctrllimited: Optional[bool] = None + gear: Optional[float] = None + + def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: + if root is None: + motor = ET.Element("motor") + else: + motor = ET.SubElement(root, "motor") + if self.name is not None: + motor.set("name", self.name) + if self.joint is not None: + motor.set("joint", self.joint) + if self.ctrllimited is not None: + motor.set("ctrllimited", str(self.ctrllimited).lower()) + if self.ctrlrange is not None: + motor.set("ctrlrange", " ".join(map(str, self.ctrlrange))) + if self.gear is not None: + motor.set("gear", str(self.gear)) + return motor + + +@dataclass +class Light: + directional: bool = True + diffuse: Optional[Tuple[float, float, float]] = None + specular: Optional[Tuple[float, float, float]] = None + pos: Optional[Tuple[float, float, float]] = None + dir: Optional[Tuple[float, float, float]] = None + castshadow: Optional[bool] = None + + def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: + if root is None: + light = ET.Element("light") + else: + light = ET.SubElement(root, "light") + if self.directional is not None: + light.set("directional", str(self.directional).lower()) + if self.diffuse is not None: + light.set("diffuse", " ".join(map(str, self.diffuse))) + if self.specular is not None: + light.set("specular", " ".join(map(str, self.specular))) + if self.pos is not None: + light.set("pos", " ".join(map(str, self.pos))) + if self.dir is not None: + light.set("dir", " ".join(map(str, self.dir))) + if self.castshadow is not None: + light.set("castshadow", str(self.castshadow).lower()) + return light + + +@dataclass +class Equality: + solref: Tuple[float, float] + + def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: + equality = ET.Element("equality") if root is None else ET.SubElement(root, "equality") + equality.set("solref", " ".join(map(str, self.solref))) + return equality + + +@dataclass +class Site: + name: Optional[str] = None + size: Optional[float] = None + pos: Optional[Tuple[float, float, float]] = None + + def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: + site = ET.Element("site") if root is None else ET.SubElement(root, "site") + if self.name is not None: + site.set("name", self.name) + if self.size is not None: + site.set("size", str(self.size)) + if self.pos is not None: + site.set("pos", " ".join(map(str, self.pos))) + return site + + +@dataclass +class Default: + joint: Optional[Joint] = None + geom: Optional[Geom] = None + class_: Optional[str] = None + motor: Optional[Motor] = None + equality: Optional[Equality] = None + + def to_xml(self, default: Optional[ET.Element] = None) -> ET.Element: + default = ET.Element("default") if default is None else ET.SubElement(default, "default") + if self.joint is not None: + self.joint.to_xml(default) + if self.geom is not None: + self.geom.to_xml(default) + if self.class_ is not None: + default.set("class", self.class_) + if self.motor is not None: + self.motor.to_xml(default) + if self.equality is not None: + self.equality.to_xml(default) + return default + + +@dataclass +class Actuator: + motors: List[Motor] + + def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: + actuator = ET.Element("actuator") if root is None else ET.SubElement(root, "actuator") + for motor in self.motors: + motor.to_xml(actuator) + return actuator + + +@dataclass +class Actuatorpos: + name: Optional[str] = None + actuator: Optional[str] = None + user: Optional[str] = None + + def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: + actuatorpos = ET.Element("actuatorpos") if root is None else ET.SubElement(root, "actuatorpos") + if self.name is not None: + actuatorpos.set("name", self.name) + if self.actuator is not None: + actuatorpos.set("actuator", self.actuator) + if self.user is not None: + actuatorpos.set("user", self.user) + return actuatorpos + + +@dataclass +class Actuatorvel: + name: Optional[str] = None + actuator: Optional[str] = None + user: Optional[str] = None + + def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: + actuatorvel = ET.Element("actuatorvel") if root is None else ET.SubElement(root, "actuatorvel") + if self.name is not None: + actuatorvel.set("name", self.name) + if self.actuator is not None: + actuatorvel.set("actuator", self.actuator) + if self.user is not None: + actuatorvel.set("user", self.user) + return actuatorvel + + +@dataclass +class Actuatorfrc: + name: Optional[str] = None + actuator: Optional[str] = None + user: Optional[str] = None + noise: Optional[float] = None + + def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: + actuatorfrc = ET.Element("actuatorfrc") if root is None else ET.SubElement(root, "actuatorfrc") + if self.name is not None: + actuatorfrc.set("name", self.name) + if self.actuator is not None: + actuatorfrc.set("actuator", self.actuator) + if self.user is not None: + actuatorfrc.set("user", self.user) + if self.noise is not None: + actuatorfrc.set("noise", str(self.noise)) + return actuatorfrc + + +@dataclass +class Sensor: + actuatorpos: Optional[List[Actuatorpos]] = None + actuatorvel: Optional[List[Actuatorvel]] = None + actuatorfrc: Optional[List[Actuatorfrc]] = None + + def to_xml(self, root: Optional[ET.Element] = None) -> ET.Element: + sensor = ET.Element("sensor") if root is None else ET.SubElement(root, "sensor") + if self.actuatorpos is not None: + for actuatorpos in self.actuatorpos: + actuatorpos.to_xml(sensor) + if self.actuatorvel is not None: + for actuatorvel in self.actuatorvel: + actuatorvel.to_xml(sensor) + if self.actuatorfrc is not None: + for actuatorfrc in self.actuatorfrc: + actuatorfrc.to_xml(sensor) + return sensor + + +def _copy_stl_files(source_directory: Union[str, Path], destination_directory: Union[str, Path]) -> None: + # Ensure the destination directory exists, create if not + os.makedirs(destination_directory, exist_ok=True) + + # Use glob to find all .stl files in the source directory + pattern = os.path.join(source_directory, "*.stl") + for file_path in glob.glob(pattern): + destination_path = os.path.join(destination_directory, os.path.basename(file_path)) + shutil.copy(file_path, destination_path) + print(f"Copied {file_path} to {destination_path}") + + +def _remove_stl_files(source_directory: Union[str, Path]) -> None: + for filename in os.listdir(source_directory): + if filename.endswith(".stl"): + file_path = os.path.join(source_directory, filename) + os.remove(file_path) + + +class Robot: + """A class to adapt the world in a Mujoco XML file.""" + + def __init__( + self, + robot_name: str, + output_dir: Union[str, Path], + compiler: Optional[Compiler] = None, + remove_inertia: bool = True, + ) -> None: + """Initialize the robot. + + Args: + robot_name (str): The name of the robot. + output_dir (str | Path): The output directory. + compiler (Compiler, optional): The compiler settings. + remove_inertia (bool, optional): Whether to remove inertia tags. + """ + self.robot_name = robot_name + self.output_dir = output_dir + self.compiler = compiler + if remove_inertia: + self._set_clean_up() + self.tree = ET.parse(self.output_dir / f"{self.robot_name}.xml") + + def _set_clean_up(self) -> None: + # HACK + # mujoco has a hard time reading meshes + _copy_stl_files(self.output_dir / "meshes", self.output_dir) + # remove inertia tags + urdf_tree = ET.parse(self.output_dir / f"{self.robot_name}.urdf") + root = urdf_tree.getroot() + for link in root.findall(".//link"): + inertial = link.find("inertial") + if inertial is not None: + link.remove(inertial) + + tree = ET.ElementTree(root) + tree.write(self.output_dir / f"{self.robot_name}.urdf", encoding="utf-8", xml_declaration=True) + model = mujoco.MjModel.from_xml_path((self.output_dir / f"{self.robot_name}.urdf").as_posix()) + mujoco.mj_saveLastXML((self.output_dir / f"{self.robot_name}.xml").as_posix(), model) + # remove all the files + _remove_stl_files(self.output_dir) + + def adapt_world(self) -> None: + root = self.tree.getroot() + + compiler = root.find("compiler") + if self.compiler is not None: + compiler = self.compiler.to_xml(compiler) + + worldbody = root.find("worldbody") + new_root_body = Body(name="root", pos=(0, 0, 0), quat=(1, 0, 0, 0)).to_xml() + + # List to store items to be moved to the new root body + items_to_move = [] + # Gather all children (geoms and bodies) that need to be moved under the new root body + for element in worldbody: + items_to_move.append(element) + # 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) + self.tree = ET.ElementTree(root) + + def save(self, path: Union[str, Path]) -> None: + self.tree.write(path, encoding="utf-8", xml_declaration=True) diff --git a/sim/scripts/create_mjcf.py b/sim/scripts/create_mjcf.py index 91c95bec..1ff62dd0 100644 --- a/sim/scripts/create_mjcf.py +++ b/sim/scripts/create_mjcf.py @@ -15,9 +15,9 @@ import xml.dom.minidom import xml.etree.ElementTree as ET from pathlib import Path +from typing import List, Union -from kol.formats import mjcf - +from sim.formats import mjcf from sim.stompy.joints import StompyFixed logger = logging.getLogger(__name__) @@ -38,23 +38,23 @@ def adapt_world(self) -> None: root: ET.Element = self.tree.getroot() root.append( - mjcf.Option( # type: ignore[call-arg] + mjcf.Option( timestep=0.001, viscosity=1e-6, iterations=50, solver="PGS", gravity=(0, 0, -9.81), - flag=mjcf.Flag(frictionloss="enable"), # type: ignore[attr-defined] + flag=mjcf.Flag(frictionloss="enable"), ).to_xml() ) # TODO - test the physical parameters root.append( - mjcf.Default( # type: ignore[attr-defined] - joint=mjcf.Joint(armature=0.01, damping=0.1, limited=True, frictionloss=0.01), # type: ignore[call-arg] - motor=mjcf.Motor(ctrllimited=True), # type: ignore[attr-defined] - equality=mjcf.Equality(solref=(0.001, 2)), # type: ignore[attr-defined] - geom=mjcf.Geom( # type: ignore[call-arg] + mjcf.Default( + joint=mjcf.Joint(armature=0.01, damping=0.1, limited=True, frictionloss=0.01), + motor=mjcf.Motor(ctrllimited=True), + equality=mjcf.Equality(solref=(0.001, 2)), + geom=mjcf.Geom( solref=(0.001, 2), friction=(0.9, 0.2, 0.2), condim=4, @@ -100,16 +100,16 @@ def adapt_world(self) -> None: # Add joints to all the movement of the base new_root_body.extend( [ - mjcf.Joint(name="root_x", type="slide", axis=(1, 0, 0), limited=False).to_xml(), # type: ignore[call-arg] - mjcf.Joint(name="root_y", type="slide", axis=(0, 1, 0), limited=False).to_xml(), # type: ignore[call-arg] - mjcf.Joint(name="root_z", type="slide", axis=(0, 0, 1), limited=False).to_xml(), # type: ignore[call-arg] - mjcf.Joint(name="root_ball", type="ball", limited=False).to_xml(), # type: ignore[call-arg] + mjcf.Joint(name="root_x", type="slide", axis=(1, 0, 0), limited=False).to_xml(), + mjcf.Joint(name="root_y", type="slide", axis=(0, 1, 0), limited=False).to_xml(), + mjcf.Joint(name="root_z", type="slide", axis=(0, 0, 1), limited=False).to_xml(), + mjcf.Joint(name="root_ball", type="ball", limited=False).to_xml(), ] ) # Add imu site to the body - relative position to the body # check at what stage we use this - new_root_body.append(mjcf.Site(name="imu", size=0.01, pos=(0, 0, 0)).to_xml()) # type: ignore[attr-defined] + 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: @@ -121,7 +121,7 @@ def adapt_world(self) -> None: worldbody.insert( 0, - mjcf.Light( # type: ignore[attr-defined] + mjcf.Light( directional=True, diffuse=(0.4, 0.4, 0.4), specular=(0.1, 0.1, 0.1), @@ -132,13 +132,13 @@ def adapt_world(self) -> None: ) worldbody.insert( 0, - mjcf.Light( # type: ignore[attr-defined] + mjcf.Light( directional=True, diffuse=(0.6, 0.6, 0.6), specular=(0.2, 0.2, 0.2), pos=(0, 0, 4), dir=(0, 0, -1) ).to_xml(), ) worldbody.insert( 0, - mjcf.Geom( # type: ignore[call-arg] + mjcf.Geom( name="ground", type="plane", size=(0, 0, 1), @@ -150,26 +150,24 @@ def adapt_world(self) -> None: ).to_xml(), ) - motors = [] - sensors = [] + motors: List[mjcf.Motor] = [] + sensor_pos: List[mjcf.Actuatorpos] = [] + sensor_vel: List[mjcf.Actuatorvel] = [] + sensor_frc: List[mjcf.Actuatorfrc] = [] for joint, limits in StompyFixed.default_limits().items(): if joint in StompyFixed.legs.all_joints(): motors.append( - mjcf.Motor( # type: ignore[attr-defined] + mjcf.Motor( name=joint, joint=joint, gear=1, ctrlrange=(limits["lower"], limits["upper"]), ctrllimited=True ) ) - sensors.extend( - [ - mjcf.Actuatorpos(name=joint + "_p", actuator=joint, user="13"), # type: ignore[attr-defined] - mjcf.Actuatorvel(name=joint + "_v", actuator=joint, user="13"), # type: ignore[attr-defined] - mjcf.Actuatorfrc(name=joint + "_f", actuator=joint, user="13", noise=0.001), # type: ignore[attr-defined] - ] - ) + sensor_pos.append(mjcf.Actuatorpos(name=joint + "_p", actuator=joint, user="13")) + sensor_vel.append(mjcf.Actuatorvel(name=joint + "_v", actuator=joint, user="13")) + sensor_frc.append(mjcf.Actuatorfrc(name=joint + "_f", actuator=joint, user="13", noise=0.001)) # Add motors and sensors - root.append(mjcf.Actuator(motors).to_xml()) # type: ignore[arg-type, call-arg] - root.append(mjcf.Sensor(sensors).to_xml()) # type: ignore[attr-defined] + root.append(mjcf.Actuator(motors).to_xml()) + root.append(mjcf.Sensor(sensor_pos, sensor_vel, sensor_frc).to_xml()) # Add imus sensors = root.find("sensor") @@ -186,7 +184,7 @@ def adapt_world(self) -> None: ) self.tree = ET.ElementTree(root) - def save(self, path: str | Path) -> None: + def save(self, path: Union[str, Path]) -> None: rough_string = ET.tostring(self.tree.getroot(), "utf-8") # Pretty print the XML formatted_xml = _pretty_print_xml(rough_string) diff --git a/sim/scripts/simulate_mjcf.py b/sim/scripts/simulate_mjcf.py index 9efde71b..921e6d1d 100644 --- a/sim/scripts/simulate_mjcf.py +++ b/sim/scripts/simulate_mjcf.py @@ -9,6 +9,7 @@ import logging import time from pathlib import Path +from typing import List, Union import mediapy as media import mujoco @@ -21,8 +22,8 @@ logger = logging.getLogger(__name__) -def simulate(model_path: str | Path, duration: float, framerate: float, record_video: bool) -> None: - frames: list[np.ndarray] = [] +def simulate(model_path: Union[str, Path], duration: float, framerate: float, record_video: bool) -> None: + frames: List[np.ndarray] = [] model = mujoco.MjModel.from_xml_path(model_path) data = mujoco.MjData(model) renderer = mujoco.Renderer(model) diff --git a/sim/scripts/update_stompy_s3.py b/sim/scripts/update_stompy_s3.py index c19f5f49..b214010d 100644 --- a/sim/scripts/update_stompy_s3.py +++ b/sim/scripts/update_stompy_s3.py @@ -1,11 +1,9 @@ +# mypy: disable-error-code="import-not-found" """Updates the Stompy model.""" import tarfile from pathlib import Path -from kol.logging import configure_logging -from kol.onshape.converter import Converter - STOMPY_MODEL = ( "https://cad.onshape.com/documents/71f793a23ab7562fb9dec82d/" "w/6160a4f44eb6113d3fa116cd/e/1a95e260677a2d2d5a3b1eb3" @@ -23,6 +21,12 @@ def main() -> None: + try: + from kol.logging import configure_logging + from kol.onshape.converter import Converter + except ImportError: + raise ImportError("Please install the `kscale-onshape-library` package to run this script.") + configure_logging() output_dir = Path("stompy") From 6080de96d76b2c0383acb9e99be6cd710c30a633 Mon Sep 17 00:00:00 2001 From: Ben Bolte Date: Wed, 15 May 2024 18:59:42 -0700 Subject: [PATCH 3/6] Update README.md --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index 7c0cd362..eb581131 100755 --- a/README.md +++ b/README.md @@ -29,6 +29,8 @@ The getting up task is still an open challenge! ## Getting Started +This repository requires Python 3.8 due to compatibility issues with underlying libraries. We hope to support more recent Python versions in the future. + 1. Clone this repository: ```bash git clone https://github.com/kscalelabs/sim.git From fb0bbb2c1e728e2ba255c7355fbd2a548e475253 Mon Sep 17 00:00:00 2001 From: Ben Bolte Date: Wed, 15 May 2024 19:02:02 -0700 Subject: [PATCH 4/6] Update README.md --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index eb581131..7019678a 100755 --- a/README.md +++ b/README.md @@ -10,6 +10,7 @@ [![Discord](https://img.shields.io/discord/1224056091017478166)](https://discord.gg/k5mSvCkYQh) [![Wiki](https://img.shields.io/badge/wiki-humanoids-black)](https://humanoids.wiki)
+[![Python Checks](https://github.com/kscalelabs/sim/actions/workflows/test.yml/badge.svg)](https://github.com/kscalelabs/sim/actions/workflows/test.yml) [![Update Stompy S3 Model](https://github.com/kscalelabs/sim/actions/workflows/update_stompy_s3.yml/badge.svg)](https://github.com/kscalelabs/sim/actions/workflows/update_stompy_s3.yml) From 27509783198426024668a8936680a6c757b28ed8 Mon Sep 17 00:00:00 2001 From: Benjamin Bolte Date: Wed, 15 May 2024 19:02:40 -0700 Subject: [PATCH 5/6] more --- tests/conftest.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index fdc62bec..26b37caa 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,7 @@ """Defines PyTest configuration for the project.""" import random +from typing import List import pytest from _pytest.python import Function @@ -11,5 +12,5 @@ def set_random_seed() -> None: random.seed(1337) -def pytest_collection_modifyitems(items: list[Function]) -> None: +def pytest_collection_modifyitems(items: List[Function]) -> None: items.sort(key=lambda x: x.get_closest_marker("slow") is not None) From 963f09410ec6b64e00b004f4fd99902d04c837fe Mon Sep 17 00:00:00 2001 From: Ben Bolte Date: Wed, 15 May 2024 19:10:14 -0700 Subject: [PATCH 6/6] Update README.md --- README.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/README.md b/README.md index 7019678a..613fbae9 100755 --- a/README.md +++ b/README.md @@ -10,6 +10,10 @@ [![Discord](https://img.shields.io/discord/1224056091017478166)](https://discord.gg/k5mSvCkYQh) [![Wiki](https://img.shields.io/badge/wiki-humanoids-black)](https://humanoids.wiki)
+[![python](https://img.shields.io/badge/-Python_3.8-blue?logo=python&logoColor=white)](https://github.com/pre-commit/pre-commit) +[![black](https://img.shields.io/badge/Code%20Style-Black-black.svg?labelColor=gray)](https://black.readthedocs.io/en/stable/) +[![ruff](https://img.shields.io/badge/Linter-Ruff-red.svg?labelColor=gray)](https://github.com/charliermarsh/ruff) +
[![Python Checks](https://github.com/kscalelabs/sim/actions/workflows/test.yml/badge.svg)](https://github.com/kscalelabs/sim/actions/workflows/test.yml) [![Update Stompy S3 Model](https://github.com/kscalelabs/sim/actions/workflows/update_stompy_s3.yml/badge.svg)](https://github.com/kscalelabs/sim/actions/workflows/update_stompy_s3.yml)