VictorMylle/RL-Sim-Framework
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

✨ better robot joint loading

Browse Source
This commit is contained in:
Victor Mylle
2026-03-09 22:17:28 +01:00
parent 9be07d9186
commit 70cd2cdd7d
13 changed files with 215 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files

83
src/runners/mujoco.py
View File

@@ -6,7 +6,8 @@ import mujoco
import numpy as np
import torch
from src.core.env import BaseEnv, ActuatorConfig
from src.core.env import BaseEnv
from src.core.robot import RobotConfig
from src.core.runner import BaseRunner, BaseRunnerConfig
@dataclasses.dataclass
@@ -30,18 +31,18 @@ class MuJoCoRunner(BaseRunner[MuJoCoRunnerConfig]):
return torch.device(self.config.device)
@staticmethod
def _load_model_with_actuators(model_path: str, actuators: list[ActuatorConfig]) -> mujoco.MjModel:
"""Load a URDF (or MJCF) file and programmatically inject actuators.
def _load_model(robot: RobotConfig) -> mujoco.MjModel:
"""Load a URDF (or MJCF) and apply robot.yaml settings.
Two-step approach required because MuJoCo's URDF parser ignores
<actuator> in the <mujoco> extension block:
1. Load the URDF → MuJoCo converts it to internal MJCF
2. Export the MJCF XML, add <actuator> elements, reload
2. Export the MJCF XML, inject actuators + joint overrides, reload
This keeps the URDF clean and standard — actuator config lives in
the env config (Isaac Lab pattern), not in the robot file.
This keeps the URDF clean (re-exportable from CAD) — all hardware
tuning lives in robot.yaml.
"""
abs_path = Path(model_path).resolve()
abs_path = robot.urdf_path.resolve()
model_dir = abs_path.parent
is_urdf = abs_path.suffix.lower() == ".urdf"
@@ -74,33 +75,45 @@ class MuJoCoRunner(BaseRunner[MuJoCoRunnerConfig]):
else:
model_raw = mujoco.MjModel.from_xml_path(str(abs_path))
if not actuators:
if not robot.actuators and not robot.joints:
return model_raw
# Step 2: Export internal MJCF representation (save next to original
# model so relative mesh/asset paths resolve correctly on reload)
# Step 2: Export internal MJCF, inject actuators + joint overrides, reload
tmp_mjcf = model_dir / "_tmp_actuator_inject.xml"
try:
mujoco.mj_saveLastXML(str(tmp_mjcf), model_raw)
mjcf_str = tmp_mjcf.read_text()
# Step 3: Inject actuators into the MJCF XML
# Use torque actuator. Speed is limited by joint damping:
# at steady state, vel_max = torque / damping.
root = ET.fromstring(mjcf_str)
act_elem = ET.SubElement(root, "actuator")
for act in actuators:
ET.SubElement(act_elem, "motor", attrib={
"name": f"{act.joint}_motor",
"joint": act.joint,
"gear": str(act.gear),
"ctrlrange": f"{act.ctrl_range[0]} {act.ctrl_range[1]}",
})
# Add damping to actuated joints to limit max speed and
# mimic real motor friction / back-EMF.
# vel_max ≈ max_torque / damping
joint_damping = {a.joint: a.damping for a in actuators}
# ── Inject actuators ────────────────────────────────────
if robot.actuators:
act_elem = ET.SubElement(root, "actuator")
for act in robot.actuators:
attribs = {
"name": f"{act.joint}_{act.type}",
"joint": act.joint,
"gear": str(act.gear),
"ctrlrange": f"{act.ctrl_range[0]} {act.ctrl_range[1]}",
}
if act.type == "position":
attribs["kp"] = str(act.kp)
if act.kv > 0:
attribs["kv"] = str(act.kv)
ET.SubElement(act_elem, "position", attrib=attribs)
elif act.type == "velocity":
attribs["kp"] = str(act.kp)
ET.SubElement(act_elem, "velocity", attrib=attribs)
else: # motor (default)
ET.SubElement(act_elem, "motor", attrib=attribs)
# ── Apply joint overrides from robot.yaml ───────────────
# Merge actuator damping + explicit joint overrides
joint_damping = {a.joint: a.damping for a in robot.actuators}
for name, jcfg in robot.joints.items():
if jcfg.damping is not None:
joint_damping[name] = jcfg.damping
for body in root.iter("body"):
for jnt in body.findall("joint"):
name = jnt.get("name")
@@ -115,6 +128,15 @@ class MuJoCoRunner(BaseRunner[MuJoCoRunnerConfig]):
geom.set("contype", "0")
geom.set("conaffinity", "0")
# Harden joint limits: MuJoCo's default soft limits are too
# weak and allow overshoot. Negative solref = hard constraint
# (direct stiffness/damping instead of impedance match).
for body in root.iter("body"):
for jnt in body.findall("joint"):
if jnt.get("limited") == "true" or jnt.get("range"):
jnt.set("solreflimit", "-1000 -100")
jnt.set("solimplimit", "0.95 0.99 0.001")
# Step 4: Write modified MJCF and reload from file path
# (from_xml_path resolves mesh paths relative to the file location)
modified_xml = ET.tostring(root, encoding="unicode")
@@ -124,12 +146,7 @@ class MuJoCoRunner(BaseRunner[MuJoCoRunnerConfig]):
tmp_mjcf.unlink(missing_ok=True)
def _sim_initialize(self, config: MuJoCoRunnerConfig) -> None:
model_path = self.env.config.model_path
if model_path is None:
raise ValueError("model_path must be specified in the environment config")
actuators = self.env.config.actuators
self._model = self._load_model_with_actuators(str(model_path), actuators)
self._model = self._load_model(self.env.robot)
self._model.opt.timestep = config.dt
self._data: list[mujoco.MjData] = [mujoco.MjData(self._model) for _ in range(config.num_envs)]
@@ -195,10 +212,6 @@ class MuJoCoRunner(BaseRunner[MuJoCoRunnerConfig]):
torch.from_numpy(qvel_batch).to(self.device),
)
def _sim_close(self) -> None:
if hasattr(self, "_offscreen_renderer") and self._offscreen_renderer is not None:
self._offscreen_renderer.close()
def render(self, env_idx: int = 0) -> np.ndarray:
"""Offscreen render of a single environment."""
if not hasattr(self, "_offscreen_renderer"):