tau-ctrl 0.1.0__py3-none-any.whl

tau_ctrl/__init__.py

@@ -0,0 +1,75 @@
+"""
+tau-ctrl — an SB3-like, simulator-agnostic controller/algorithm framework.
+
+Call algorithms directly off the package, the way you would with SB3:
+
+    from tau_ctrl import make, PID, MPPI, CEM, ICEM, ILQR, CBFFilter, PPO, SAC, TD3
+
+    ctrl = make("mppi", env, horizon=25, n_samples=200)   # or PID(env), PPO(env), ...
+    action, _ = ctrl.predict(obs)                          # SB3-style
+    ctrl.learn(total_timesteps=100_000)                    # trainable methods (ppo/sac/td3)
+
+Every algorithm shares one interface (``predict``/``learn``/``save``/``load``)
+over a plain ``gymnasium.Env`` — feedback (PID), sampling-based MPC (MPPI,
+CEM, ICEM), gradient-based MPC (ILQR), a safety filter (CBF), on-policy RL
+(PPO), and off-policy RL (SAC, TD3 — replay-buffer based, far more
+sample-efficient than PPO for continuous control).
+
+Model-based methods (MPPI/CEM/ICEM/ILQR/CBF) need the env to be *branchable*
+— to expose ``get_state()``/``set_state()``. Feedback and RL work on any gym
+env. :mod:`tau_ctrl.tuning` adds Bayesian/genetic auto-tuning of any gains.
+"""
+
+__version__ = "0.1.0"
+
+from .algorithms import (
+    CBFFilter,
+    CEM,
+    ICEM,
+    ILQR,
+    MPPI,
+    PID,
+    PPO,
+    SAC,
+    TD3,
+    BaseController,
+    ControllerManifest,
+    Trainer,
+    available,
+    get_state,
+    is_branchable,
+    make,
+    probe_env,
+    probe_hardware,
+    register,
+    resolve_device,
+    set_state,
+)
+from .tuning import AutoTuner
+
+__all__ = [
+    "make",
+    "register",
+    "available",
+    "is_branchable",
+    "get_state",
+    "set_state",
+    "resolve_device",
+    "BaseController",
+    "ControllerManifest",
+    "PID",
+    "MPPI",
+    "CEM",
+    "ICEM",
+    "ILQR",
+    "CBFFilter",
+    "PPO",
+    "SAC",
+    "TD3",
+    "Trainer",
+    "probe_env",
+    "probe_hardware",
+    "AutoTuner",
+    "algorithms",
+    "tuning",
+]

tau_ctrl/algorithms/__init__.py

@@ -0,0 +1,91 @@
+"""tau-ctrl's SB3-like, simulator-agnostic controller/algorithm framework.
+
+One interface for the whole controller spectrum over a Gymnasium env, called
+the way you'd call SB3:
+
+    from tau_ctrl import make, PID, MPPI, CEM, ICEM, ILQR, CBFFilter, PPO, SAC, TD3
+
+    ctrl = make("mppi", env, horizon=25, n_samples=200)   # or PID(env), PPO(env), ...
+    action, _ = ctrl.predict(obs)                          # SB3-style
+    ctrl.learn(total_timesteps=100_000)                    # for trainable ones (ppo/sac/td3)
+
+Design goals:
+- **Simulator-agnostic:** talks only to the gym API (reset/step/spaces). No
+  MuJoCo/MJX assumptions — works with whatever simulator hands it an env.
+- **GPU-optional:** torch device is auto-selected for learned methods; sampling
+  planners parallelise across a vectorized env when one is supplied.
+- **Unified:** feedback (PID), sampling-MPC (MPPI/CEM/ICEM), gradient-based
+  MPC (ILQR), safety (CBF), and RL (PPO on-policy; SAC/TD3 off-policy) all
+  share ``predict``/``learn``/``save``/``load``.
+
+Model-based methods (MPPI/CEM/ICEM/ILQR/CBF) need the env to be *branchable*
+— to expose ``get_state()``/``set_state()``. Feedback and RL work on any gym
+env.
+"""
+
+from .base import (
+    BaseController,
+    BranchingNotSupported,
+    ControllerManifest,
+    available,
+    get_state,
+    is_branchable,
+    make,
+    register,
+    resolve_device,
+    set_state,
+)
+from .cbf import CBFFilter
+from .ilqr import ILQR
+from .mppi import CEM, ICEM, MPPI
+from .pid import PID
+from .ppo import PPO
+from .sac import SAC
+from .adapters import GymVectorAdapter, SyncTorchVecEnv, jax_to_torch
+from .strategy import (
+    EnvCaps,
+    HardwareCaps,
+    Plan,
+    Strategy,
+    Trainer,
+    probe_env,
+    probe_hardware,
+    select_strategy,
+)
+from .td3 import TD3
+from .vec_env import TorchPendulum, TorchVecEnv
+
+__all__ = [
+    "BaseController",
+    "ControllerManifest",
+    "make",
+    "register",
+    "available",
+    "is_branchable",
+    "get_state",
+    "set_state",
+    "BranchingNotSupported",
+    "resolve_device",
+    "PID",
+    "MPPI",
+    "CEM",
+    "ICEM",
+    "ILQR",
+    "CBFFilter",
+    "PPO",
+    "SAC",
+    "TD3",
+    "TorchVecEnv",
+    "TorchPendulum",
+    "GymVectorAdapter",
+    "SyncTorchVecEnv",
+    "jax_to_torch",
+    "Trainer",
+    "Strategy",
+    "Plan",
+    "EnvCaps",
+    "HardwareCaps",
+    "probe_env",
+    "probe_hardware",
+    "select_strategy",
+]

tau_ctrl/algorithms/adapters.py

@@ -0,0 +1,150 @@
+"""Adapters: make *whatever env you actually have* speak the TorchVecEnv
+contract, so the same SAC/TD3 vectorized loop runs unchanged.
+
+You rarely own the environment. The batched, tensor-native
+:class:`~tau_ctrl.algorithms.vec_env.TorchVecEnv` is the ideal, but real
+projects hand you one of three things, and tau-ctrl adapts to all three behind
+the same trainer (the loop only needs ``num_envs`` + ``reset`` + ``step`` →
+``(next_obs, reward, done, start_obs)``):
+
+1. **Already batched & on-device** — MJX / Brax (JAX) or Isaac Gym (torch).
+   Wrap with :func:`jax_to_torch` (zero-copy via dlpack) or, for a torch-native
+   batched env, nothing at all — it already fits the contract.
+
+2. **Batched but numpy/CPU** — a ``gymnasium.vector.VectorEnv``.
+   :class:`GymVectorAdapter` moves each step's batch to the device with one
+   copy. Collection stays CPU-bound, but the replay buffer and gradient update
+   are batched and on-device.
+
+3. **A single, non-batchable env** — PyBullet, classic MuJoCo, a real robot.
+   You can't turn its dynamics into tensor ops, but you can run ``N`` copies.
+   :class:`SyncTorchVecEnv` steps them in a Python loop and presents the batch
+   on-device. No GPU *env* parallelism (the env is the bottleneck), but the
+   exact same trainer runs and the update still saturates the GPU.
+
+If you truly have one env and can't make copies (a single real robot), just use
+the ordinary single-env ``learn()`` path — there the GPU only accelerates the
+update, which is honest: you can't parallelize physics you don't control.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+import numpy as np
+
+from .vec_env import TorchVecEnv
+
+
+def jax_to_torch(x, device: str):
+    """Zero-copy JAX array → torch tensor via dlpack (for MJX / Brax on GPU).
+
+    Both frameworks expose ``__dlpack__``; the buffer is shared, so no host
+    round-trip and no extra device allocation. Use this inside a thin
+    ``TorchVecEnv`` subclass that drives an MJX/Brax ``step`` and hands the
+    resulting device arrays straight to the trainer.
+    """
+    import torch  # noqa: PLC0415
+
+    t = torch.from_dlpack(x)
+    return t.to(device) if str(t.device) != str(device) else t
+
+
+class GymVectorAdapter(TorchVecEnv):
+    """Wrap a ``gymnasium.vector.VectorEnv`` (numpy batches) as a TorchVecEnv.
+
+    Each ``step`` moves the batch to the device with a single H→D copy. The env
+    still runs on CPU (that's inherent to gym's vector API), but the buffer and
+    update are batched on-device — the same win as a native ``TorchVecEnv``
+    minus the on-device env stepping.
+
+    Handles gym autoreset: on the step where an env is done, gym returns the
+    *reset* observation and stashes the true terminal obs in
+    ``info["final_observation"]``. We store the terminal obs as ``next_obs``
+    (correct bootstrap target) and carry the reset obs as ``start_obs``.
+    """
+
+    def __init__(self, venv: Any, device: str = "cpu") -> None:
+        super().__init__(getattr(venv, "num_envs"), device)
+        self.venv = venv
+        # gym VectorEnv exposes single_* spaces describing one env.
+        self.observation_space = getattr(venv, "single_observation_space", None) or venv.observation_space
+        self.action_space = getattr(venv, "single_action_space", None) or venv.action_space
+
+    def _to_t(self, arr):
+        t = self._torch
+        return t.as_tensor(np.asarray(arr, dtype=np.float32), device=self.device)
+
+    def reset(self):
+        obs, _ = self.venv.reset()
+        return self._to_t(obs)
+
+    def step(self, actions):
+        t = self._torch
+        a = actions.detach().cpu().numpy()
+        obs, reward, terminated, truncated, info = self.venv.step(a)
+        done = np.asarray(terminated, dtype=np.float32)
+        start_obs = np.asarray(obs, dtype=np.float32).copy()
+        next_obs = start_obs.copy()  # true post-dynamics obs
+
+        # Recover terminal observations for envs that just reset.
+        finals = info.get("final_observation") if isinstance(info, dict) else None
+        if finals is not None:
+            for i, fo in enumerate(finals):
+                if fo is not None:
+                    next_obs[i] = np.asarray(fo, dtype=np.float32)
+        return (
+            self._to_t(next_obs),
+            self._to_t(reward),
+            t.as_tensor(done, device=self.device),
+            self._to_t(start_obs),
+        )
+
+
+class SyncTorchVecEnv(TorchVecEnv):
+    """Run ``N`` copies of a single (non-batchable) gym env as one on-device batch.
+
+    For PyBullet, classic MuJoCo, or any env whose dynamics you can't rewrite as
+    tensor ops but can instantiate ``N`` times. The envs step in a plain Python
+    loop on CPU — so you get no GPU env parallelism — but the same vectorized
+    SAC/TD3 loop runs unchanged and the replay buffer + gradient update stay
+    batched on-device.
+
+    ``env_fns`` is a list of zero-arg factories (one per env), matching the SB3
+    / gym.vector convention so seeds/configs can differ per env.
+    """
+
+    def __init__(self, env_fns: list[Callable[[], Any]], device: str = "cpu") -> None:
+        self.envs = [fn() for fn in env_fns]
+        super().__init__(len(self.envs), device)
+        self.observation_space = self.envs[0].observation_space
+        self.action_space = self.envs[0].action_space
+
+    def reset(self):
+        obs = np.stack([np.asarray(e.reset()[0], dtype=np.float32) for e in self.envs])
+        return self._torch.as_tensor(obs, device=self.device)
+
+    def step(self, actions):
+        t = self._torch
+        a = actions.detach().cpu().numpy()
+        n = self.num_envs
+        next_obs = np.empty((n, *self.observation_space.shape), dtype=np.float32)
+        start_obs = np.empty_like(next_obs)
+        reward = np.empty(n, dtype=np.float32)
+        done = np.zeros(n, dtype=np.float32)
+        for i, e in enumerate(self.envs):
+            o, r, term, trunc, _ = e.step(a[i])
+            next_obs[i] = np.asarray(o, dtype=np.float32)
+            reward[i] = r
+            done[i] = float(term)  # truncation resets but doesn't break bootstrap
+            if term or trunc:
+                o2, _ = e.reset()
+                start_obs[i] = np.asarray(o2, dtype=np.float32)
+            else:
+                start_obs[i] = next_obs[i]
+        return (
+            t.as_tensor(next_obs, device=self.device),
+            t.as_tensor(reward, device=self.device),
+            t.as_tensor(done, device=self.device),
+            t.as_tensor(start_obs, device=self.device),
+        )

tau_ctrl/algorithms/base.py

@@ -0,0 +1,207 @@
+"""Core abstractions for tau-ctrl's algorithm framework — an SB3-like,
+simulator-agnostic controller framework over a Gymnasium environment.
+
+Every method (PID, MPPI, CEM-MPC, CBF safety filter, PPO, ...) implements the
+same :class:`BaseController` interface — ``predict(obs) -> action`` like SB3,
+plus an optional ``learn()`` for the ones that train. Nothing here imports a
+simulator: controllers talk only to the ``gymnasium.Env`` API. Methods that
+need to roll the dynamics forward (MPPI, CEM, CBF) do so through the optional
+*branching* capability (:func:`get_state` / :func:`set_state`), which any env
+can provide; nothing is MuJoCo- or MJX-specific.
+"""
+
+from __future__ import annotations
+
+import pickle
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Optional
+
+import numpy as np
+
+# ---------------------------------------------------------------------------
+# Registry
+# ---------------------------------------------------------------------------
+
+_REGISTRY: dict[str, type["BaseController"]] = {}
+
+
+def register(name: str):
+    def deco(cls):
+        cls.name = name
+        _REGISTRY[name] = cls
+        return cls
+
+    return deco
+
+
+def make(name: str, env: Any, **kwargs: Any) -> "BaseController":
+    """Construct a controller by name, SB3-style: ``make("mppi", env, ...)``."""
+    if name not in _REGISTRY:
+        raise KeyError(f"unknown controller {name!r}; available: {sorted(_REGISTRY)}")
+    return _REGISTRY[name](env, **kwargs)
+
+
+def available() -> list[str]:
+    return sorted(_REGISTRY)
+
+
+# ---------------------------------------------------------------------------
+# Manifest
+# ---------------------------------------------------------------------------
+
+@dataclass
+class ControllerManifest:
+    name: str
+    family: str                     # "feedback" | "sampling_mpc" | "safety" | "rl"
+    requires_branching: bool = False  # needs get_state/set_state for rollouts
+    uses_reward: bool = False          # uses the env's reward as its objective
+    trainable: bool = False            # has a meaningful learn()
+    gpu_capable: bool = False
+    extra: dict = field(default_factory=dict)
+
+
+# ---------------------------------------------------------------------------
+# Branching capability (simulator-agnostic)
+# ---------------------------------------------------------------------------
+
+class BranchingNotSupported(RuntimeError):
+    pass
+
+
+def get_state(env: Any) -> Any:
+    """Snapshot the env's dynamical state so a rollout can branch from it.
+
+    Resolution order (no simulator assumptions): an explicit ``get_state`` on
+    the (unwrapped) env, else classic-control ``.state``. Simulators that want
+    model-based control just expose ``get_state``/``set_state``.
+    """
+    u = getattr(env, "unwrapped", env)
+    if hasattr(u, "get_state"):
+        return u.get_state()
+    if hasattr(u, "state") and u.state is not None:
+        return np.array(u.state, dtype=float)
+    raise BranchingNotSupported(
+        f"{type(u).__name__} exposes no get_state()/.state; model-based methods "
+        "(MPPI/CEM/CBF) need a branchable env."
+    )
+
+
+def set_state(env: Any, state: Any) -> None:
+    u = getattr(env, "unwrapped", env)
+    if hasattr(u, "set_state"):
+        u.set_state(state)
+        return
+    if hasattr(u, "state"):
+        u.state = np.array(state, dtype=float)
+        return
+    raise BranchingNotSupported(f"{type(u).__name__} exposes no set_state().")
+
+
+def is_branchable(env: Any) -> bool:
+    try:
+        get_state(env)
+        return True
+    except BranchingNotSupported:
+        return False
+
+
+# ---------------------------------------------------------------------------
+# Device (only relevant for the learned/torch methods)
+# ---------------------------------------------------------------------------
+
+def resolve_device(pref: str = "auto") -> str:
+    if pref not in ("auto", "cuda", "cpu"):
+        return pref
+    if pref == "cpu":
+        return "cpu"
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            return "cuda"
+    except Exception:
+        pass
+    return "cpu"
+
+
+# ---------------------------------------------------------------------------
+# Base controller
+# ---------------------------------------------------------------------------
+
+class BaseController(ABC):
+    """SB3-like controller: ``predict`` always; ``learn`` when trainable."""
+
+    name: str = "base"
+    manifest: ControllerManifest
+
+    def __init__(self, env: Any, device: str = "auto", seed: Optional[int] = None) -> None:
+        self.env = env
+        self.observation_space = getattr(env, "observation_space", None)
+        self.action_space = getattr(env, "action_space", None)
+        self.device = resolve_device(device)
+        self.seed = seed
+        self.rng = np.random.default_rng(seed)
+        if seed is not None:
+            # Seed torch's global RNG too, before _setup() builds any network —
+            # weight init and stochastic-policy sampling (torch.randn_like) use
+            # torch's own RNG, not self.rng, so this is required for
+            # reproducibility in any torch-based controller (PPO/SAC/TD3/...).
+            try:
+                import torch
+
+                torch.manual_seed(seed)
+            except ImportError:
+                pass
+        self._setup()
+
+    # Subclasses override _setup / predict; learn is optional.
+    def _setup(self) -> None:  # noqa: B027
+        pass
+
+    @abstractmethod
+    def predict(
+        self, obs: Any, state: Any = None, deterministic: bool = True
+    ) -> tuple[np.ndarray, Any]:
+        """Return ``(action, next_internal_state)`` — SB3 signature."""
+
+    def learn(self, total_timesteps: int = 0, **kwargs: Any) -> "BaseController":
+        """Train the controller. Analytic controllers are no-ops."""
+        return self
+
+    def reset(self) -> None:
+        """Reset any internal controller state (call at episode start)."""
+
+    # --- helpers ---------------------------------------------------------
+    def _clip_action(self, a: np.ndarray) -> np.ndarray:
+        a = np.asarray(a, dtype=float)
+        sp = self.action_space
+        if sp is not None and getattr(sp, "low", None) is not None:
+            a = np.clip(a, sp.low, sp.high)
+        return a
+
+    @property
+    def action_dim(self) -> int:
+        return int(np.prod(self.action_space.shape))
+
+    # --- persistence -----------------------------------------------------
+    def _state_dict(self) -> dict:
+        """Override to persist trained parameters."""
+        return {}
+
+    def _load_state_dict(self, sd: dict) -> None:  # noqa: B027
+        pass
+
+    def save(self, path: str | Path) -> None:
+        blob = {"name": self.name, "state": self._state_dict()}
+        with open(path, "wb") as f:
+            pickle.dump(blob, f)
+
+    @classmethod
+    def load(cls, path: str | Path, env: Any, **kwargs: Any) -> "BaseController":
+        with open(path, "rb") as f:
+            blob = pickle.load(f)
+        obj = cls(env, **kwargs)
+        obj._load_state_dict(blob.get("state", {}))
+        return obj

tau_ctrl/algorithms/cbf.py

@@ -0,0 +1,98 @@
+"""Control Barrier Function safety filter (discrete-time, model-free).
+
+Wraps any base controller and minimally edits its action so a safety margin
+``h(state) >= 0`` is maintained. It enforces the discrete-time CBF condition
+
+    h(next_state) >= (1 - alpha) * h(state),   alpha in (0, 1]
+
+Since we make no simulator assumptions, the one-step effect of an action on
+``h`` is measured by *probing* the branchable env (set_state -> step -> read
+h). The barrier is linearized in the action via finite differences, giving a
+single-inequality QP whose solution is a closed-form projection of the nominal
+action. This keeps it fully simulator-agnostic and cheap.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Sequence
+
+import numpy as np
+
+from .base import BaseController, ControllerManifest, get_state, register, set_state
+
+Barrier = Callable[[Any], float]  # state -> margin (>= 0 is safe)
+
+
+@register("cbf")
+class CBFFilter(BaseController):
+    """Safety filter: ``u = argmin ||u - u_nom||^2 s.t. discrete-CBF holds``."""
+
+    manifest = ControllerManifest(
+        name="cbf", family="safety", requires_branching=True,
+    )
+
+    def __init__(
+        self,
+        env: Any,
+        base: BaseController,
+        barriers: Barrier | Sequence[Barrier],
+        alpha: float = 0.5,
+        fd_eps: float = 1e-3,
+        passes: int = 3,
+        **kw: Any,
+    ) -> None:
+        self.base = base
+        self.barriers: list[Barrier] = list(barriers) if isinstance(barriers, (list, tuple)) else [barriers]
+        self.alpha = float(alpha)
+        self.fd_eps = float(fd_eps)
+        self.passes = int(passes)
+        super().__init__(env, **kw)
+
+    def reset(self) -> None:
+        self.base.reset()
+
+    def _h_next(self, state0: Any, u: np.ndarray, h: Barrier) -> float:
+        set_state(self.env, state0)
+        self.env.step(u)
+        val = h(get_state(self.env))
+        return float(val)
+
+    def predict(self, obs, state=None, deterministic: bool = True):
+        u_nom, _ = self.base.predict(obs, state, deterministic)
+        u = np.array(u_nom, dtype=float).ravel()
+        state0 = get_state(self.env)
+
+        for _ in range(self.passes):
+            adjusted = False
+            for h in self.barriers:
+                h0 = float(h(state0))
+                margin = self._h_next(state0, u, h) - (1.0 - self.alpha) * h0
+                if margin >= 0:
+                    continue
+                # Linearize h_next in u via finite differences: g = ∂h_next/∂u.
+                # Perturb *inward* per dimension so that when u sits on an
+                # actuator bound the probe isn't clipped away (which would
+                # falsely read a zero gradient).
+                g = np.zeros_like(u)
+                base_val = self._h_next(state0, u, h)
+                sp = self.action_space
+                hi = getattr(sp, "high", None)
+                for i in range(u.size):
+                    step = self.fd_eps
+                    if hi is not None and u[i] + step > hi[i]:
+                        step = -self.fd_eps
+                    up = u.copy()
+                    up[i] += step
+                    g[i] = (self._h_next(state0, up, h) - base_val) / step
+                denom = float(g @ g)
+                if denom < 1e-12:
+                    continue  # action has no first-order effect on this barrier
+                # Project: min ||Δ||^2 s.t. gᵀΔ + margin >= 0  ->  Δ = -margin g/‖g‖²
+                u = u + (-margin) * g / denom
+                u = self._clip_action(u)
+                adjusted = True
+            if not adjusted:
+                break
+
+        set_state(self.env, state0)  # restore the real env
+        return self._clip_action(u), None

tau_ctrl/algorithms/envs/__init__.py

@@ -0,0 +1,6 @@
+"""Pure-Python toy envs (no simulator dependency) for tests and examples."""
+
+from ..vec_env import TorchPendulum, TorchVecEnv
+from .toy import DoubleIntegrator, PendulumSwingUp
+
+__all__ = ["PendulumSwingUp", "DoubleIntegrator", "TorchVecEnv", "TorchPendulum"]