Source code for state_space_setup_assistant.ros2_control_export

"""Build the C++-consumable ``<name>_ros2_control.yaml`` artifact.

``controller.npz`` (written by ``ControllerResult.save_npz``) carries the
plant/controller matrices but not ``joint_names``/``q_eq`` -- those live only
on the linearized ``StateSpaceModel``. A ros2_control controller needs both
(to turn a raw measured joint position into the deviation state the gain/
compensator was designed for), so this module builds one flat dict from the
model *and* the result together, meant to be dumped with
``yamlgen.dump_yaml`` right after ``result.save_npz`` in ``export.py``.

Schema (flat nested lists only, so a C++ ``yaml-cpp`` reader needs no NumPy
support):

    schema_version: 1
    controller_type: lqr            # or lqg / hinf / hinf_mixsyn / pid
    name: lqr
    joint_names: [...]              # full state order: x = [q - q_eq, qdot]
    q_eq: [...]
    actuated_joint_names: [...]     # order matches K's rows / u
    u_eq: [...]
    output_names: [...]             # used only by dynamic-compensator designs
    K: [[...], ...]                 # present iff a static-gain design
    ctrl_A/ctrl_B/ctrl_C/ctrl_D: [[...], ...]   # present instead of K otherwise
    info: {...}                     # informational scalars only, not consumed
"""

from typing import Dict

import numpy as np


def _tolist(arr) -> list:
    return np.asarray(arr, dtype=float).tolist()


def _info_dict(info: Dict) -> Dict:
    out = {}
    for key, val in info.items():
        if not np.isscalar(val):
            continue
        try:
            out[key] = float(val)
        except TypeError:
            continue  # e.g. complex scalars: not yaml-safe, informational only
    return out




def ros2_control_yaml_dict(model, result) -> Dict:
    """Build the ``*_ros2_control.yaml`` dict from a ``StateSpaceModel``
    (``session.model``) and a ``ControllerResult`` (``session.result``).

    Exactly one of ``result.K`` / ``result.controller`` must be set, matching
    ``ControllerResult``'s own invariant.
    """
    if result.K is None and result.controller is None:
        raise ValueError(
            'result has neither a static gain nor a dynamic controller')

    cfg = {
        'schema_version': 1,
        'controller_type': result.name,
        'name': result.name,
        'joint_names': list(model.joint_names),
        'q_eq': _tolist(model.q_eq),
        'actuated_joint_names': list(model.actuated_joint_names),
        'u_eq': _tolist(model.u_eq),
        'output_names': list(model.output_names),
    }

    if result.K is not None:
        cfg['K'] = _tolist(result.K)
    if result.controller is not None:
        k = result.controller
        cfg['ctrl_A'] = _tolist(k.A)
        cfg['ctrl_B'] = _tolist(k.B)
        cfg['ctrl_C'] = _tolist(k.C)
        cfg['ctrl_D'] = _tolist(k.D)

    info = _info_dict(result.info)
    if info:
        cfg['info'] = info

    return cfg