virtualmodelcontrol 0.1.0__py3-none-any.whl

virtualmodelcontrol/models/continuum/pcc.py

@@ -0,0 +1,120 @@
+"""Piecewise constant curvature (PCC) kinematics of a continuum robot with n segments."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import casadi as ca
+import numpy as np
+
+from ...core.params import ParamSet, as_param
+from ...core.registry import register
+from ...core.space import Euclidean
+from ...core.units import M
+
+
+def segment_frame(dx: Any, dy: Any, dl: Any, L0: Any, d: Any, s: Any, eps: float) -> tuple:
+    """(R, t) of one PCC segment at local arc fraction s ∈ [0, 1].
+
+    The segment is an arc of radius d (L0 + Dl) / D, D = √(Dx² + Dy² + ε), bent by s D / d
+    about the axis (−Dy, Dx, 0) / D. Its tangent is R's z axis.
+    """
+    D = ca.sqrt(dx**2 + dy**2 + eps)
+    theta = s * D / d
+    c, sn = ca.cos(theta), ca.sin(theta)
+    R = ca.vertcat(
+        ca.horzcat(1 + dx**2 / D**2 * (c - 1), dx * dy / D**2 * (c - 1), dx / D * sn),
+        ca.horzcat(dx * dy / D**2 * (c - 1), 1 + dy**2 / D**2 * (c - 1), dy / D * sn),
+        ca.horzcat(-dx / D * sn, -dy / D * sn, c),
+    )
+    k = d * (L0 + dl) / D**2
+    t = ca.vertcat(k * dx * (1 - c), k * dy * (1 - c), k * D * sn)
+    return R, t
+
+
+@register("model", "pcc")
+class PCC:
+    """Continuum robot of n PCC segments; q = Δ = [Dx, Dy, Dl] per segment, base to tip [m].
+
+    Per segment i (from 1): ``seg{i}.L0`` rest length and ``seg{i}.d`` section radius [m],
+    ``design`` Params. The arc parameter s ∈ [0, 1] is uniform in arc length (the breakpoints
+    follow from the rest lengths). Sites: ``base``, ``seg{i}`` (end of segment i), ``tip``.
+    """
+
+    q_unit = M
+
+    def __init__(self, L0: Sequence[Any], d: Any, *, eps: float = 1e-12) -> None:
+        n = len(L0)
+        d_list = list(d) if isinstance(d, (list, tuple, np.ndarray)) else [d] * n
+        if len(d_list) != n:
+            raise ValueError(f"{n} rest lengths but {len(d_list)} section radii")
+        self.n_segments = n
+        self.space = Euclidean(3 * n)
+        self.eps = eps
+        self.params = ParamSet()
+        for i in range(n):
+            for name, value in (("L0", L0[i]), ("d", d_list[i])):
+                key = f"seg{i + 1}.{name}"
+                param = as_param(value, key, unit=M, bounds=(0.0, np.inf), scope="design")
+                self.params.add(param, key)
+        self.sites = ("base", *(f"seg{i + 1}" for i in range(n)), "tip")
+
+    def frame(self, q: Any, at: Any, p: dict[str, Any]) -> tuple[Any, Any]:
+        """(R, position) at a site or at arc parameter s (clipped to [0, 1])."""
+        n = self.n_segments
+        L0 = [p[f"seg{i + 1}.L0"] for i in range(n)]
+        d = [p[f"seg{i + 1}.d"] for i in range(n)]
+        seg = [q[3 * i : 3 * i + 3] for i in range(n)]
+
+        def local(i: int, s: Any) -> tuple:
+            return segment_frame(seg[i][0], seg[i][1], seg[i][2], L0[i], d[i], s, self.eps)
+
+        if isinstance(at, str):
+            if at not in self.sites:
+                raise KeyError(f"unknown site {at!r}; sites: {self.sites}")
+            count = n if at == "tip" else 0 if at == "base" else int(at[3:])
+            R, pos = ca.DM.eye(3), ca.DM.zeros(3, 1)
+            for i in range(count):
+                Ri, ti = local(i, 1.0)
+                R, pos = ca.mtimes(R, Ri), pos + ca.mtimes(R, ti)
+            return R, pos
+
+        cum = [0]
+        for i in range(n):
+            cum.append(cum[-1] + L0[i])
+        b = [c / cum[-1] for c in cum]  # breakpoints of s at the segment ends
+        s = ca.fmin(ca.fmax(at, 0.0), 1.0)
+        R_base, p_base = ca.DM.eye(3), ca.DM.zeros(3, 1)
+        candidates = []
+        for i in range(n):
+            Ri, ti = local(i, (s - b[i]) / (b[i + 1] - b[i]))
+            candidates.append((ca.mtimes(R_base, Ri), p_base + ca.mtimes(R_base, ti)))
+            Re, te = local(i, 1.0)
+            R_base, p_base = ca.mtimes(R_base, Re), p_base + ca.mtimes(R_base, te)
+        R, pos = candidates[-1]
+        for i in reversed(range(n - 1)):
+            inside = s <= b[i + 1]
+            R = ca.if_else(inside, candidates[i][0], R)
+            pos = ca.if_else(inside, candidates[i][1], pos)
+        return R, pos
+
+    def breakpoints(self) -> np.ndarray:
+        """Arc parameter at the base and at each segment end, at the current rest lengths."""
+        L0 = np.array([float(self.params[f"seg{i + 1}.L0"].value) for i in range(self.n_segments)])
+        return np.concatenate([[0.0], np.cumsum(L0)]) / L0.sum()
+
+    def to_dict(self) -> dict[str, Any]:
+        """Constructor arguments at the current Param values."""
+        n = range(1, self.n_segments + 1)
+        return {
+            "type": "pcc",
+            "L0": [float(self.params[f"seg{i}.L0"].value) for i in n],
+            "d": [float(self.params[f"seg{i}.d"].value) for i in n],
+            "eps": self.eps,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> PCC:
+        """Inverse of ``to_dict``."""
+        return cls(data["L0"], data["d"], eps=data.get("eps", 1e-12))

virtualmodelcontrol/models/kinematic.py

@@ -0,0 +1,41 @@
+"""Kinematic model contract: a space, Params, named sites and frames written in CasADi."""
+
+from __future__ import annotations
+
+from typing import Any, Protocol
+
+import casadi as ca
+import numpy as np
+from numpy.typing import ArrayLike
+
+from ..core.params import ParamSet, constants
+from ..core.registry import get
+from ..core.space import Space
+
+
+class KinematicModel(Protocol):
+    """Frames of a robot as smooth functions of q and the Params.
+
+    ``frame(q, at, p)`` returns (R, position) in the base frame; ``at`` is a site name or, for a
+    continuous body, an arc parameter s (possibly symbolic); ``p`` maps the model's Param names
+    to CasADi expressions.
+    """
+
+    space: Space
+    params: ParamSet
+    sites: tuple[str, ...]
+
+    def frame(self, q: Any, at: Any, p: dict[str, Any]) -> tuple[Any, Any]:
+        """Rotation (3, 3) and position (3, 1) at ``at``."""
+        ...
+
+
+def evaluate_frame(model: KinematicModel, q: ArrayLike, at: Any) -> tuple[np.ndarray, np.ndarray]:
+    """Numeric (R, position) at the model's current Param values."""
+    R, p = model.frame(ca.DM(np.asarray(q, dtype=float)), at, constants(model.params))
+    return np.array(ca.evalf(R)), np.array(ca.evalf(p)).ravel()
+
+
+def from_dict(data: dict[str, Any], kind: str = "model") -> Any:
+    """Build a registered model (or actuation, with ``kind``) from ``to_dict()`` output."""
+    return get(kind, data["type"]).from_dict(data)

virtualmodelcontrol/models/rigid/__init__.py

@@ -0,0 +1,6 @@
+"""Rigid models."""
+
+from .couplings import LinearCoupling
+from .poe import SerialChain
+
+__all__ = ["LinearCoupling", "SerialChain"]

virtualmodelcontrol/models/rigid/couplings.py

@@ -0,0 +1,64 @@
+"""Couplings: a model whose joints are driven together, joint angles = matrix · q."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import casadi as ca
+import numpy as np
+
+from ...core.params import Param, ParamSet
+from ...core.registry import get, register
+from ...core.space import Euclidean
+
+
+@register("model", "linear_coupling")
+class LinearCoupling:
+    """Wraps a joint-space model so that its joint angles are ``coupling`` · q.
+
+    Use it when one motor drives several joints (pulleys, mimic joints): q holds one entry per
+    motor. ``coupling`` (n_joints × n_q) is a ``design`` Param; the wrapped model's Params keep
+    their names.
+    """
+
+    def __init__(self, model: Any, coupling: Any) -> None:
+        value = np.atleast_2d(np.asarray(getattr(coupling, "value", coupling), dtype=float))
+        if value.shape[0] != model.space.nq:
+            raise ValueError(
+                f"coupling has {value.shape[0]} rows but the model has {model.space.nq} joints"
+            )
+        self.model = model
+        self.space = Euclidean(value.shape[1])
+        self.params = ParamSet()
+        self.params.merge(model.params)
+        free = (-np.inf, np.inf)
+        self.coupling = (
+            coupling
+            if isinstance(coupling, Param)
+            else Param("coupling", value, scope="design", bounds=free)
+        )
+        self.params.add(self.coupling, "coupling")
+        self.sites = tuple(model.sites)
+        self.q_unit = getattr(model, "q_unit", "")
+
+    def joint_angles(self, q: Any, p: dict[str, Any]) -> Any:
+        """Joint angles of the wrapped model."""
+        return ca.mtimes(p["coupling"], q)
+
+    def frame(self, q: Any, at: Any, p: dict[str, Any]) -> tuple[Any, Any]:
+        """Frame of the wrapped model at the coupled joint angles."""
+        return self.model.frame(self.joint_angles(q, p), at, p)
+
+    def to_dict(self) -> dict[str, Any]:
+        """The wrapped model and the coupling matrix."""
+        return {
+            "type": "linear_coupling",
+            "model": self.model.to_dict(),
+            "coupling": self.coupling.value.tolist(),
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> LinearCoupling:
+        """Inverse of ``to_dict``."""
+        inner = data["model"]
+        return cls(get("model", inner["type"]).from_dict(inner), data["coupling"])

virtualmodelcontrol/models/rigid/poe.py

@@ -0,0 +1,95 @@
+"""Serial chains by the product of exponentials, with all geometry as Params."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+import casadi as ca
+import numpy as np
+
+from ...core.params import ParamSet, as_param
+from ...core.registry import register
+from ...core.space import Euclidean
+from ...core.symbolic import exp_so3
+from ...core.units import M
+
+JOINT_TYPES = ("revolute", "prismatic")
+
+
+@register("model", "poe")
+class SerialChain:
+    """Serial chain of revolute and prismatic joints, as a product of exponentials.
+
+    Joint i (from 1) has ``j{i}.axis`` (direction at q = 0, normalized internally) and, if revolute,
+    ``j{i}.point`` [m], a point on its axis. A site has ``{name}.position`` [m] at q = 0 and moves
+    with the joints up to ``after`` (1-based). All are ``design`` Params.
+    """
+
+    def __init__(
+        self,
+        joints: Sequence[str],
+        axes: Sequence[Any],
+        points: Sequence[Any],
+        sites: Mapping[str, tuple[int, Any]],
+    ) -> None:
+        if any(j not in JOINT_TYPES for j in joints):
+            raise ValueError(f"joint types must be in {JOINT_TYPES}, got {list(joints)}")
+        if not len(joints) == len(axes) == len(points):
+            raise ValueError("give one axis and one point per joint")
+        self.joints = list(joints)
+        self.space = Euclidean(len(joints))
+        self.params = ParamSet()
+        free = (-np.inf, np.inf)
+        for i, (axis, point) in enumerate(zip(axes, points, strict=True)):
+            key = f"j{i + 1}.axis"
+            self.params.add(as_param(axis, key, scope="design", bounds=free), key)
+            key = f"j{i + 1}.point"
+            self.params.add(as_param(point, key, unit=M, scope="design", bounds=free), key)
+        self.site_joint: dict[str, int] = {}
+        for name, (after, position) in sites.items():
+            key = f"{name}.position"
+            self.params.add(as_param(position, key, unit=M, scope="design", bounds=free), key)
+            self.site_joint[name] = int(after)
+        self.sites = tuple(sites)
+
+    @property
+    def q_unit(self) -> str:
+        """Unit of the joint coordinates (rad if any joint is revolute)."""
+        return "rad" if "revolute" in self.joints else M
+
+    def frame(self, q: Any, at: str, p: dict[str, Any]) -> tuple[Any, Any]:
+        """(R, position) of a site."""
+        if at not in self.site_joint:
+            raise KeyError(f"unknown site {at!r}; sites: {self.sites}")
+        R, pos = ca.DM.eye(3), ca.DM.zeros(3, 1)
+        for i in range(self.site_joint[at]):
+            axis = ca.reshape(p[f"j{i + 1}.axis"], 3, 1)
+            axis = axis / ca.norm_2(axis)
+            if self.joints[i] == "revolute":
+                Ri = exp_so3(axis, q[i])
+                ti = ca.mtimes(ca.DM.eye(3) - Ri, ca.reshape(p[f"j{i + 1}.point"], 3, 1))
+            else:
+                Ri, ti = ca.DM.eye(3), axis * q[i]
+            R, pos = ca.mtimes(R, Ri), pos + ca.mtimes(R, ti)
+        return R, pos + ca.mtimes(R, ca.reshape(p[f"{at}.position"], 3, 1))
+
+    def to_dict(self) -> dict[str, Any]:
+        """Constructor arguments at the current Param values."""
+        n = range(1, len(self.joints) + 1)
+        return {
+            "type": "poe",
+            "joints": list(self.joints),
+            "axes": [self.params[f"j{i}.axis"].value.tolist() for i in n],
+            "points": [self.params[f"j{i}.point"].value.tolist() for i in n],
+            "sites": {
+                name: [after, self.params[f"{name}.position"].value.tolist()]
+                for name, after in self.site_joint.items()
+            },
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> SerialChain:
+        """Inverse of ``to_dict``."""
+        sites = {name: (after, pos) for name, (after, pos) in data["sites"].items()}
+        return cls(data["joints"], data["axes"], data["points"], sites)

virtualmodelcontrol/robots/__init__.py

@@ -0,0 +1,5 @@
+"""Robot data: model builders and default parameters for the lab's robots."""
+
+from . import adapt, helyx
+
+__all__ = ["adapt", "helyx"]

virtualmodelcontrol/robots/adapt.py

@@ -0,0 +1,98 @@
+"""ADAPT finger: three phalanges driven by two motors, the last two joints moving together."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from ..core.params import Param
+from ..mechanisms import Custom, Joint, LimitSpring, Mechanism, PointMass
+from ..models import LinearCoupling, SerialChain
+
+LINK_LENGTHS = (0.040, 0.030, 0.0175)
+"""Proximal, middle and distal phalanx [m], from the MCP joint to the fingertip."""
+
+MOTOR_RADIUS = 0.005  # [m] motor pulley
+FINGER_PULLEY_RADIUS = 0.00223  # [m] pulley on the MCP joint
+PIP_TRANSMISSION = 0.00813  # [m] cable transmission constant of the PIP/DIP drive
+
+COUPLING = np.array([
+    [FINGER_PULLEY_RADIUS / MOTOR_RADIUS, 0.0],
+    [0.0, MOTOR_RADIUS / PIP_TRANSMISSION],
+    [0.0, MOTOR_RADIUS / PIP_TRANSMISSION],
+])  # fmt: skip
+"""Joint angles (MCP, PIP, DIP) = COUPLING · motor angles; the DIP joint mimics the PIP joint."""
+
+LINK_MASSES = (0.0057, 0.0040, 0.025)
+"""Masses of the three phalanges [kg]."""
+
+LINK_COGS = (
+    np.array([-0.000639, 0.020, 0.001276]),
+    np.array([0.000925, 0.015, 0.001138]),
+    np.array([-0.000716, 0.0105, 0.000955]),
+)
+"""Centres of gravity in each phalanx's frame [m]."""
+
+GRAVITY = (0.0, 0.0, 9.81)
+"""Gravity in the finger's base frame [m/s²], as mounted (the frame's z axis points down)."""
+
+JOINT_LIMITS = {"MCP": (0.0, np.pi / 2), "PIP": (0.0, 1.22173), "DIP": (0.0, 1.22173)}
+"""Joint ranges [rad]."""
+
+LIMIT_STIFFNESS = 1.0  # [N·m/rad], stiffness of the joint-limit springs
+
+MOTOR_EFFICIENCY = (0.8373, 0.3594)
+"""Delivered over commanded torque of the two motors (MCP, PIP)."""
+
+TORQUE_LIMIT = 0.8  # [N·m] per motor, for an optional clip at the hardware boundary
+
+
+def finger(name: str = "finger", gravity: Any = None) -> Mechanism:
+    """The finger as a robot mechanism; q holds the two motor angles [rad].
+
+    Sites: ``pip``, ``dip``, ``tip`` and the phalanges' centres of gravity ``*_cog``.
+    """
+    a, b, c = LINK_LENGTHS
+    x = [1.0, 0.0, 0.0]
+    joints = ([0.0, 0.0, 0.0], [0.0, a, 0.0], [0.0, a + b, 0.0])
+    sites = {
+        "pip": (1, joints[1]),
+        "dip": (2, joints[2]),
+        "tip": (3, [0.0, a + b + c, 0.0]),
+    }
+    for i, link in enumerate(("mcp", "pip", "dip")):
+        sites[f"{link}_cog"] = (i + 1, np.asarray(joints[i]) + LINK_COGS[i])
+    model = LinearCoupling(SerialChain(["revolute"] * 3, [x, x, x], list(joints), sites), COUPLING)
+    robot = Mechanism(name, model=model)
+    robot.add_param(
+        Param(
+            "gravity",
+            GRAVITY if gravity is None else gravity,
+            unit="m/s^2",
+            scope="design",
+            bounds=(-np.inf, np.inf),
+        )
+    )
+    for i, link in enumerate(("mcp", "pip", "dip")):
+        robot.add(f"m_{link}", PointMass(robot.point(f"{link}_cog"), LINK_MASSES[i]))
+    return robot
+
+
+def joint_angles(robot: Mechanism) -> Custom:
+    """Coordinate of the three joint angles (MCP, PIP, DIP) [rad] of a finger."""
+    coupling = robot.model.coupling
+    return Custom(
+        lambda motors, coupling: coupling @ motors,
+        [Joint(slice(0, 2), unit="rad")],
+        dim=3,
+        unit="rad",
+        params={"coupling": coupling},
+    )
+
+
+def joint_limit_spring(robot: Mechanism, stiffness: Any = LIMIT_STIFFNESS) -> LimitSpring:
+    """A spring that pushes each joint back inside its range (zero force inside)."""
+    lower = [lim[0] for lim in JOINT_LIMITS.values()]
+    upper = [lim[1] for lim in JOINT_LIMITS.values()]
+    return LimitSpring(joint_angles(robot), stiffness, lower, upper)

virtualmodelcontrol/robots/helyx.py

@@ -0,0 +1,79 @@
+"""Helyx soft arm: three tendon-driven PCC segments, in the lab's three geometries."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from ..core.params import Param
+from ..mechanisms import Gravity, LinearDamper, LinearSpring, Mechanism, PointMass
+from ..models import PCC, TendonTransmission
+
+GEOMETRIES: dict[str, dict[str, Any]] = {
+    "145-145-145": {"L0": (0.145, 0.145, 0.145), "gravity": (0.0, -9.81, 0.0)},
+    "145-290-290": {"L0": (0.145, 0.290, 0.290), "gravity": (0.0, 0.0, 9.81)},
+    "290-145-145": {"L0": (0.290, 0.145, 0.145), "gravity": (0.0, 0.0, -9.81)},
+}
+"""Segment rest lengths base to tip [m], and gravity in the base frame as mounted [m/s²]:
+side-mounted, hanging, and pointing up."""
+
+ENCODER_SIGN: dict[str, float] = {"145-145-145": -1.0, "145-290-290": -1.0, "290-145-145": 1.0}
+"""Sign of each arm's motor encoders against the library convention (θ > 0 pulls a tendon)."""
+
+SECTION_RADIUS = 0.030  # [m]
+SPOOL_RADIUS = 0.003  # [m]
+TENDON_ANGLES = np.radians([[0.0, 120.0, -120.0], [60.0, 180.0, -60.0], [150.0, 270.0, 30.0]])
+"""Tendon angles around the section at each segment's base [rad]."""
+SEGMENT_MASS = 0.040  # [kg] per REFERENCE_LENGTH of segment, lumped at the segment's midpoint
+REFERENCE_LENGTH = 0.145  # [m]
+
+
+def arm(geometry: str = "145-290-290", name: str = "arm", gravity: Any = None) -> Mechanism:
+    """The arm as a robot mechanism: PCC model, tendons, lumped masses and a gravity Param.
+
+    ``gravity`` [m/s², base frame] overrides the geometry's mounting.
+    """
+    spec = GEOMETRIES[geometry]
+    model = PCC(spec["L0"], SECTION_RADIUS)
+    robot = Mechanism(
+        name, model=model, actuation=TendonTransmission(list(TENDON_ANGLES), SPOOL_RADIUS)
+    )
+    g = spec["gravity"] if gravity is None else gravity
+    robot.add_param(Param("gravity", g, unit="m/s^2", scope="design", bounds=(-np.inf, np.inf)))
+    b = model.breakpoints()
+    for i, length in enumerate(spec["L0"]):
+        mass = SEGMENT_MASS * (length / REFERENCE_LENGTH)
+        robot.add(f"m{i + 1}", PointMass(robot.point(s=(b[i] + b[i + 1]) / 2), mass))
+    return robot
+
+
+SIM_STIFFNESS = np.array([
+    233.9329569634883, 199.3719645292524, 410.00362883032545,
+    391.6936839908778, 475.436633573179, 831.8383997691384,
+    445.34472425862594, 517.8961860162584, 938.3740743393556,
+]) / 0.12  # fmt: skip
+"""Diagonal stiffness in Δ of the simulated 145-290-290 arm [N/m], referred to commanded torque
+(identified values referred to delivered torque, divided by the efficiency 0.12)."""
+
+SIM_DAMPING = np.array([
+    125.13438533964865, 104.5563494024772, 205.43182586195965,
+    139.28436710395044, 148.1129490938869, 284.33900221754124,
+    135.3243281016147, 141.08179561149416, 280.464403371687,
+]) / 0.12  # fmt: skip
+"""Diagonal damping in Δ of the simulated 145-290-290 arm [N·s/m], referred like the stiffness."""
+
+
+def add_dynamics(robot: Mechanism, stiffness: Any = None, damping: Any = None) -> Mechanism:
+    """Give the arm its physical stiffness and damping in Δ and gravity, for simulation.
+
+    ``stiffness`` [N/m] and ``damping`` [N·s/m] are per-axis (9 values); the defaults are the
+    simulated arm's.
+    """
+    delta = robot.joint(slice(0, 9))
+    K = SIM_STIFFNESS if stiffness is None else stiffness
+    D = SIM_DAMPING if damping is None else damping
+    robot.add("stiffness", LinearSpring(delta, Param("stiffness", K, unit="N/m", scope="design")))
+    robot.add("damping", LinearDamper(delta, Param("damping", D, unit="N*s/m", scope="design")))
+    robot.add("gravity", Gravity(robot))
+    return robot

virtualmodelcontrol/sim/__init__.py

@@ -0,0 +1,7 @@
+"""Simulation and the run loop: plants, a model-based simulator, guard and recorder."""
+
+from .model_plant import ModelPlant
+from .plant import Plant, SimPlant
+from .run import Guard, RunLog, SimClock, run
+
+__all__ = ["Guard", "ModelPlant", "Plant", "RunLog", "SimClock", "SimPlant", "run"]

virtualmodelcontrol/sim/model_plant.py

@@ -0,0 +1,79 @@
+"""ModelPlant: simulates a robot from its own components, and reads like the hardware."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Iterable
+from typing import Any
+
+import numpy as np
+from numpy.typing import ArrayLike
+
+from ..core.signals import Signals
+from ..dynamics import compile_dynamics
+from ..mechanisms.mechanism import Mechanism
+
+
+class ModelPlant:
+    """Simulated robot: linearly implicit Euler steps of at most ``max_step`` [s].
+
+    ``read`` returns motor angles and rates (as the hardware does) plus q and v. The command
+    ``motor_torque`` is held between calls to ``advance``.
+    """
+
+    def __init__(
+        self,
+        robot: Mechanism,
+        q0: ArrayLike | None = None,
+        v0: ArrayLike | None = None,
+        *,
+        max_step: float = 1e-3,
+        runtime: Iterable[str] = (),
+    ) -> None:
+        self.dynamics = compile_dynamics(robot, runtime)
+        self.space = robot.model.space
+        self.max_step = max_step
+        self.p = self.dynamics.live_values()
+        self._q0 = self.space.neutral() if q0 is None else np.asarray(q0, dtype=float)
+        self._v0 = np.zeros(self.space.nv) if v0 is None else np.asarray(v0, dtype=float)
+        self.reset()
+
+    def reset(self, x0: Any = None, seed: int | None = None) -> None:
+        """Back to (q0, v0), or to ``x0 = (q, v)``; time 0, zero command."""
+        q, v = (self._q0, self._v0) if x0 is None else x0
+        self.q: np.ndarray = np.array(q, dtype=float)
+        self.v: np.ndarray = np.array(v, dtype=float)
+        self.u: np.ndarray = np.zeros(self.dynamics.n_u)
+        self.t = 0.0
+
+    def read(self) -> Signals:
+        """Motor angles and rates, plus the configuration and velocity."""
+        theta, theta_dot = self.dynamics.motors(self.q, self.v, self.p)
+        return Signals(
+            self.t,
+            motor_position=np.array(theta).ravel(),
+            motor_velocity=np.array(theta_dot).ravel(),
+            q=self.q,
+            v=self.v,
+        )
+
+    def write(self, cmd: Signals) -> None:
+        """Hold ``motor_torque`` [N·m] until the next write."""
+        self.u = np.array(cmd["motor_torque"], dtype=float)
+
+    def advance(self, dt: float) -> None:
+        """Integrate over ``dt`` [s]."""
+        n = max(1, math.ceil(dt / self.max_step - 1e-9))
+        h = dt / n
+        for _ in range(n):
+            q, v = self.dynamics.step(self.q, self.v, self.u, self.p, self.t, h)
+            self.q, self.v = np.array(q).ravel(), np.array(v).ravel()
+            self.t += h
+
+    def energy(self) -> float:
+        """Kinetic plus stored energy of the robot [J]."""
+        T, V = self.dynamics.energy(self.q, self.v, self.p, self.t)
+        return float(T) + float(V)
+
+    def close(self) -> None:
+        """Nothing to release."""

virtualmodelcontrol/sim/plant.py

@@ -0,0 +1,40 @@
+"""Plants: anything that can be read and commanded, simulated or real."""
+
+from __future__ import annotations
+
+from typing import Any, Protocol
+
+from ..core.signals import Signals
+
+
+class Plant(Protocol):
+    """A robot as the controller sees it: ``read`` measurements, ``write`` commands.
+
+    Hardware and simulators implement the same three methods, so one run loop serves both.
+    """
+
+    def read(self) -> Signals:
+        """Latest measurements (``motor_position`` [rad], ``motor_velocity`` [rad/s], ...)."""
+        ...
+
+    def write(self, cmd: Signals) -> None:
+        """Apply a command (``motor_torque`` [N·m])."""
+        ...
+
+    def close(self) -> None:
+        """Release resources; a real plant stops its motors."""
+        ...
+
+
+class SimPlant(Plant, Protocol):
+    """A simulated plant: it also resets and advances its own time ``t`` [s]."""
+
+    t: float
+
+    def reset(self, x0: Any = None, seed: int | None = None) -> None:
+        """Back to an initial state."""
+        ...
+
+    def advance(self, dt: float) -> None:
+        """Integrate over ``dt`` [s] holding the last command."""
+        ...