Source code for urdf_state_space.config

"""YAML-driven model setup: describe the linearization per joint, by name.

Example config
--------------
::

    model:
      urdf: package://my_robot_description/urdf/my_robot.urdf
      floating_base: false

    joints:                        # every movable joint may have an entry
      shoulder_joint: {actuated: true, damping: 0.1}
      elbow_joint:    {actuated: true, damping: 0.1, q_eq: 0.3}
      wrist_joint:    {damping: 0.005}

    outputs:
      velocities: true             # y includes velocities as well
      # joints: [elbow_joint]      # optional; default: the actuated joints

    export:
      npz: model.npz
      mat: model.mat
      # dt: 0.001                  # also produce/export the ZOH-discretized model

If any joint sets ``actuated``, only joints with ``actuated: true`` are
actuated; if none does, all movable joints are. Joints without a ``damping``
or ``q_eq`` entry default to 0 / neutral. Relative paths (URDF and export
files) are resolved against the YAML file's directory.
"""

import os
import subprocess
from dataclasses import dataclass
from typing import Optional, Tuple

import yaml

from .state_space import StateSpaceModel, build_state_space




@dataclass
class ExportSpec:
    npz: Optional[str] = None
    mat: Optional[str] = None
    dt: Optional[float] = None





def resolve_resource(ref: str, base_dir: str = '.') -> str:
    """Resolve a package:// URI or a (possibly relative) filesystem path."""
    if ref.startswith('package://'):
        from ament_index_python.packages import get_package_share_directory
        pkg, _, rel = ref[len('package://'):].partition('/')
        return os.path.join(get_package_share_directory(pkg), rel)
    if os.path.isabs(ref):
        return ref
    return os.path.join(base_dir, ref)



def _read_urdf(path: str) -> str:
    if path.endswith('.xacro'):
        return subprocess.check_output(['xacro', path], text=True)
    with open(path) as f:
        return f.read()




def build_from_yaml(config_path: str) -> Tuple[StateSpaceModel, ExportSpec]:
    """Build a StateSpaceModel from a YAML config; also return export spec."""
    with open(config_path) as f:
        cfg = yaml.safe_load(f) or {}
    base_dir = os.path.dirname(os.path.abspath(config_path))

    model_cfg = cfg.get('model', {})
    if 'urdf' not in model_cfg:
        raise ValueError(f'{config_path}: missing required key model.urdf')
    urdf_path = resolve_resource(str(model_cfg['urdf']), base_dir)

    joints = cfg.get('joints') or {}
    for name, entry in joints.items():
        if not isinstance(entry, dict):
            raise ValueError(
                f'{config_path}: joints.{name} must be a mapping, '
                f'e.g. {{actuated: true, damping: 0.1}}')

    q_eq = {n: e['q_eq'] for n, e in joints.items() if 'q_eq' in e}
    damping = {n: e['damping'] for n, e in joints.items() if 'damping' in e}

    actuated = None
    if any('actuated' in e for e in joints.values()):
        actuated = [n for n, e in joints.items() if e.get('actuated')]
        if not actuated:
            raise ValueError(f'{config_path}: no joint has actuated: true')

    out_cfg = cfg.get('outputs') or {}
    output_joints = out_cfg.get('joints')

    model = build_state_space(
        _read_urdf(urdf_path),
        q_eq=q_eq or None,
        actuated_joints=actuated,
        output_joints=output_joints,
        velocity_outputs=bool(out_cfg.get('velocities', False)),
        joint_damping=damping or None,
        floating_base=bool(model_cfg.get('floating_base', False)),
    )

    exp = cfg.get('export') or {}
    export = ExportSpec(
        npz=resolve_resource(exp['npz'], base_dir) if 'npz' in exp else None,
        mat=resolve_resource(exp['mat'], base_dir) if 'mat' in exp else None,
        dt=float(exp['dt']) if 'dt' in exp else None,
    )
    return model, export