187 lines
6.7 KiB
Python
187 lines
6.7 KiB
Python
"""Export tuned motor parameters to MJCF and robot.yaml files.
|
|
|
|
Reads the original motor.xml and robot.yaml, patches with optimised
|
|
parameter values, and writes motor_tuned.xml + robot_tuned.yaml.
|
|
|
|
Usage:
|
|
python -m src.sysid.motor.export \
|
|
--asset-path assets/motor \
|
|
--result assets/motor/motor_sysid_result.json
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import copy
|
|
import json
|
|
import xml.etree.ElementTree as ET
|
|
from pathlib import Path
|
|
|
|
import structlog
|
|
import yaml
|
|
|
|
log = structlog.get_logger()
|
|
|
|
_DEFAULT_ASSET = "assets/motor"
|
|
|
|
|
|
def export_tuned_files(
|
|
asset_path: str | Path,
|
|
params: dict[str, float],
|
|
) -> tuple[Path, Path]:
|
|
"""Write tuned MJCF and robot.yaml files.
|
|
|
|
Returns (tuned_mjcf_path, tuned_robot_yaml_path).
|
|
"""
|
|
asset_path = Path(asset_path).resolve()
|
|
|
|
robot_yaml_path = asset_path / "robot.yaml"
|
|
robot_cfg = yaml.safe_load(robot_yaml_path.read_text())
|
|
mjcf_path = asset_path / robot_cfg["mjcf"]
|
|
|
|
# ── Tune MJCF ────────────────────────────────────────────────
|
|
tree = ET.parse(str(mjcf_path))
|
|
root = tree.getroot()
|
|
|
|
# Actuator — use average gear for the MJCF model.
|
|
gear_pos = params.get("actuator_gear_pos", params.get("actuator_gear"))
|
|
gear_neg = params.get("actuator_gear_neg", params.get("actuator_gear"))
|
|
gear_avg = None
|
|
if gear_pos is not None and gear_neg is not None:
|
|
gear_avg = (gear_pos + gear_neg) / 2.0
|
|
elif gear_pos is not None:
|
|
gear_avg = gear_pos
|
|
filter_tau = params.get("actuator_filter_tau")
|
|
for act_el in root.iter("general"):
|
|
if act_el.get("name") == "motor":
|
|
if gear_avg is not None:
|
|
act_el.set("gear", str(gear_avg))
|
|
if filter_tau is not None:
|
|
if filter_tau > 0:
|
|
act_el.set("dyntype", "filter")
|
|
act_el.set("dynprm", str(filter_tau))
|
|
else:
|
|
act_el.set("dyntype", "none")
|
|
|
|
# Joint — average damping & friction for MJCF (asymmetry in runtime).
|
|
fl_pos = params.get("motor_frictionloss_pos", params.get("motor_frictionloss"))
|
|
fl_neg = params.get("motor_frictionloss_neg", params.get("motor_frictionloss"))
|
|
fl_avg = None
|
|
if fl_pos is not None and fl_neg is not None:
|
|
fl_avg = (fl_pos + fl_neg) / 2.0
|
|
elif fl_pos is not None:
|
|
fl_avg = fl_pos
|
|
damp_pos = params.get("motor_damping_pos", params.get("motor_damping"))
|
|
damp_neg = params.get("motor_damping_neg", params.get("motor_damping"))
|
|
damp_avg = None
|
|
if damp_pos is not None and damp_neg is not None:
|
|
damp_avg = (damp_pos + damp_neg) / 2.0
|
|
elif damp_pos is not None:
|
|
damp_avg = damp_pos
|
|
for jnt in root.iter("joint"):
|
|
if jnt.get("name") == "motor_joint":
|
|
if damp_avg is not None:
|
|
jnt.set("damping", str(damp_avg))
|
|
if "motor_armature" in params:
|
|
jnt.set("armature", str(params["motor_armature"]))
|
|
if fl_avg is not None:
|
|
jnt.set("frictionloss", str(fl_avg))
|
|
|
|
# Rotor mass.
|
|
if "rotor_mass" in params:
|
|
for geom in root.iter("geom"):
|
|
if geom.get("name") == "rotor_disk":
|
|
geom.set("mass", str(params["rotor_mass"]))
|
|
|
|
# Write tuned MJCF.
|
|
tuned_mjcf_name = mjcf_path.stem + "_tuned" + mjcf_path.suffix
|
|
tuned_mjcf_path = asset_path / tuned_mjcf_name
|
|
ET.indent(tree, space=" ")
|
|
tree.write(str(tuned_mjcf_path), xml_declaration=True, encoding="unicode")
|
|
log.info("tuned_mjcf_written", path=str(tuned_mjcf_path))
|
|
|
|
# ── Tune robot.yaml ──────────────────────────────────────────
|
|
tuned_cfg = copy.deepcopy(robot_cfg)
|
|
tuned_cfg["mjcf"] = tuned_mjcf_name
|
|
|
|
if tuned_cfg.get("actuators") and len(tuned_cfg["actuators"]) > 0:
|
|
act = tuned_cfg["actuators"][0]
|
|
if gear_avg is not None:
|
|
act["gear"] = round(gear_avg, 6)
|
|
if "actuator_filter_tau" in params:
|
|
act["filter_tau"] = round(params["actuator_filter_tau"], 6)
|
|
if "motor_damping" in params:
|
|
act["damping"] = round(params["motor_damping"], 6)
|
|
|
|
if "joints" not in tuned_cfg:
|
|
tuned_cfg["joints"] = {}
|
|
if "motor_joint" not in tuned_cfg["joints"]:
|
|
tuned_cfg["joints"]["motor_joint"] = {}
|
|
mj = tuned_cfg["joints"]["motor_joint"]
|
|
if "motor_armature" in params:
|
|
mj["armature"] = round(params["motor_armature"], 6)
|
|
if fl_avg is not None:
|
|
mj["frictionloss"] = round(fl_avg, 6)
|
|
|
|
# Asymmetric / hardware-realism / nonlinear parameters.
|
|
realism = {}
|
|
for key in [
|
|
"actuator_gear_pos", "actuator_gear_neg",
|
|
"motor_damping_pos", "motor_damping_neg",
|
|
"motor_frictionloss_pos", "motor_frictionloss_neg",
|
|
"motor_deadzone_pos", "motor_deadzone_neg",
|
|
"action_bias",
|
|
"viscous_quadratic", "back_emf_gain",
|
|
"stribeck_friction_boost", "stribeck_vel",
|
|
"gearbox_backlash",
|
|
]:
|
|
if key in params:
|
|
realism[key] = round(params[key], 6)
|
|
if realism:
|
|
tuned_cfg["hardware_realism"] = realism
|
|
|
|
tuned_yaml_path = asset_path / "robot_tuned.yaml"
|
|
header = (
|
|
"# Tuned motor config — generated by src.sysid.motor.optimize\n"
|
|
"# Original: robot.yaml\n\n"
|
|
)
|
|
tuned_yaml_path.write_text(
|
|
header + yaml.dump(tuned_cfg, default_flow_style=False, sort_keys=False)
|
|
)
|
|
log.info("tuned_robot_yaml_written", path=str(tuned_yaml_path))
|
|
|
|
return tuned_mjcf_path, tuned_yaml_path
|
|
|
|
|
|
# ── CLI ──────────────────────────────────────────────────────────────
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description="Export tuned motor parameters to MJCF + robot.yaml."
|
|
)
|
|
parser.add_argument("--asset-path", type=str, default=_DEFAULT_ASSET)
|
|
parser.add_argument(
|
|
"--result", type=str, default=None,
|
|
help="Path to motor_sysid_result.json (auto-detected if omitted)",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
asset_path = Path(args.asset_path).resolve()
|
|
if args.result:
|
|
result_path = Path(args.result)
|
|
else:
|
|
result_path = asset_path / "motor_sysid_result.json"
|
|
|
|
if not result_path.exists():
|
|
raise FileNotFoundError(f"Result file not found: {result_path}")
|
|
|
|
result = json.loads(result_path.read_text())
|
|
params = result["best_params"]
|
|
|
|
export_tuned_files(asset_path=args.asset_path, params=params)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|