tau-ctrl 0.1.0__py3-none-any.whl

tau_ctrl/algorithms/envs/toy.py

@@ -0,0 +1,109 @@
+"""Tiny pure-Python Gymnasium envs for testing/examples.
+
+No simulator dependency (no MuJoCo/MJX) — they exist so tau-ctrl's algorithms
+can be tested end-to-end anywhere. Both expose ``get_state``/``set_state`` so
+the model-based controllers (MPPI/CEM/CBF) can branch rollouts, exactly as a
+real simulator env would once tau-sim adds that capability.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import numpy as np
+
+try:
+    import gymnasium as gym
+    from gymnasium import spaces
+except Exception as e:  # pragma: no cover
+    raise ImportError("tau_ctrl.algorithms requires gymnasium") from e
+
+
+class PendulumSwingUp(gym.Env):
+    """Torque-limited pendulum. angle 0 = upright; reward peaks upright.
+
+    obs = [cos(th), sin(th), thdot]. Reward = -(th^2 + 0.1*thdot^2 + 1e-3*u^2).
+    """
+
+    metadata = {"render_modes": []}
+
+    def __init__(self, max_torque: float = 6.0, dt: float = 0.05,
+                 g: float = 10.0, m: float = 1.0, l: float = 1.0,
+                 start_angle: float = np.pi, max_speed: float = 8.0):
+        self.max_torque = float(max_torque)
+        self.max_speed = float(max_speed)
+        self.dt, self.g, self.m, self.l = dt, g, m, l
+        self.start_angle = float(start_angle)
+        self.state = np.array([self.start_angle, 0.0])
+        high = np.array([1.0, 1.0, self.max_speed], dtype=np.float32)
+        self.observation_space = spaces.Box(-high, high, dtype=np.float32)
+        self.action_space = spaces.Box(-self.max_torque, self.max_torque, (1,), np.float32)
+
+    # --- branching capability ---
+    def get_state(self):
+        return np.array(self.state, dtype=float)
+
+    def set_state(self, s):
+        self.state = np.array(s, dtype=float)
+
+    def _obs(self):
+        th, thd = self.state
+        return np.array([np.cos(th), np.sin(th), thd], dtype=np.float32)
+
+    def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None):
+        super().reset(seed=seed)
+        self.state = np.array([self.start_angle, 0.0])
+        return self._obs(), {}
+
+    def step(self, action):
+        th, thd = self.state
+        u = float(np.clip(np.asarray(action).reshape(-1)[0], -self.max_torque, self.max_torque))
+        # angle measured from upright: gravity accelerates away from upright.
+        thddot = (self.g / self.l) * np.sin(th) + u / (self.m * self.l ** 2)
+        thd = np.clip(thd + thddot * self.dt, -self.max_speed, self.max_speed)
+        th = th + thd * self.dt
+        th = ((th + np.pi) % (2 * np.pi)) - np.pi  # wrap to [-pi, pi]
+        self.state = np.array([th, thd])
+        cost = th ** 2 + 0.1 * thd ** 2 + 1e-3 * u ** 2
+        return self._obs(), float(-cost), False, False, {}
+
+
+class DoubleIntegrator(gym.Env):
+    """1-D point mass: state=[x, v], action = acceleration (bounded).
+
+    Reward drives x -> 0. Handy for testing a CBF that must keep x <= x_max.
+    obs = [x, v].
+    """
+
+    metadata = {"render_modes": []}
+
+    def __init__(self, dt: float = 0.05, a_max: float = 3.0, x_max: float = 1.0,
+                 x0: float = 0.5, v0: float = 1.5):
+        self.dt, self.a_max, self.x_max = dt, float(a_max), float(x_max)
+        self.x0, self.v0 = float(x0), float(v0)
+        self.state = np.array([x0, v0])
+        self.observation_space = spaces.Box(-np.inf, np.inf, (2,), np.float32)
+        self.action_space = spaces.Box(-self.a_max, self.a_max, (1,), np.float32)
+
+    def get_state(self):
+        return np.array(self.state, dtype=float)
+
+    def set_state(self, s):
+        self.state = np.array(s, dtype=float)
+
+    def _obs(self):
+        return self.state.astype(np.float32)
+
+    def reset(self, *, seed=None, options=None):
+        super().reset(seed=seed)
+        self.state = np.array([self.x0, self.v0])
+        return self._obs(), {}
+
+    def step(self, action):
+        x, v = self.state
+        a = float(np.clip(np.asarray(action).reshape(-1)[0], -self.a_max, self.a_max))
+        v = v + a * self.dt
+        x = x + v * self.dt
+        self.state = np.array([x, v])
+        reward = -(x ** 2 + 0.1 * v ** 2)
+        return self._obs(), float(reward), False, False, {}

tau_ctrl/algorithms/ilqr.py

@@ -0,0 +1,182 @@
+"""iLQR — gradient-based trajectory optimization (a different family from
+MPPI/CEM/ICEM's sampling-based search).
+
+Where the sampling-MPC family searches by rolling out many random action
+sequences, iLQR alternates a **backward pass** (locally linearize the
+dynamics and quadratically approximate the cost along the current nominal
+trajectory, solve the resulting LQR problem for a feedback law) with a
+**forward pass** (apply that feedback law, with a backtracking line search on
+the feedforward term). It converges in far fewer trajectory evaluations when
+the dynamics are reasonably smooth, at the cost of being more sensitive to
+strong nonlinearities/discontinuities than sampling-based search.
+
+Simulator-agnostic like the rest of :mod:`tau_ctrl.algorithms`: since the env
+is an opaque ``step`` function, the dynamics Jacobians and cost derivatives
+are estimated by finite differences on a branchable env (``get_state``/
+``set_state``), the same contract MPPI/CEM/CBF already rely on. This is a
+lightweight iLQR variant: cost curvature is approximated diagonally and
+cross state-action curvature is dropped (a standard, cheap simplification —
+full analytic/autodiff Hessians would need a differentiable simulator).
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import numpy as np
+
+from .base import BaseController, ControllerManifest, get_state, register, set_state
+
+
+@register("ilqr")
+class ILQR(BaseController):
+    manifest = ControllerManifest(
+        name="ilqr", family="gradient_mpc", requires_branching=True, uses_reward=True,
+    )
+
+    def __init__(
+        self,
+        env: Any,
+        horizon: int = 20,
+        iterations: int = 5,
+        fd_eps: float = 1e-3,
+        reg: float = 1e-2,
+        line_search_steps: tuple[float, ...] = (1.0, 0.5, 0.25, 0.1, 0.0),
+        **kw: Any,
+    ) -> None:
+        self.horizon = int(horizon)
+        self.iterations = int(iterations)
+        self.fd_eps = float(fd_eps)
+        self.reg = float(reg)
+        self.line_search_steps = line_search_steps
+        super().__init__(env, **kw)
+
+    def _setup(self) -> None:
+        self.nu = self.action_dim
+        state0 = get_state(self.env)
+        self.nx = int(np.asarray(state0).size)
+        self._u_nom = np.zeros((self.horizon, self.nu))
+
+    def reset(self) -> None:
+        self._u_nom[:] = 0.0
+
+    # ------------------------------------------------------------------
+    # Finite-difference probes on the branchable env
+    # ------------------------------------------------------------------
+    def _probe(self, x: np.ndarray, u: np.ndarray) -> tuple[np.ndarray, float]:
+        """One step from (x, u): returns (next_state, reward)."""
+        set_state(self.env, x)
+        u = self._clip_action(u)
+        _, r, _, _, _ = self.env.step(u)
+        return np.asarray(get_state(self.env), dtype=float), float(r)
+
+    def _linearize(self, x: np.ndarray, u: np.ndarray):
+        """Central-difference A = df/dx, B = df/du, and cost gradient l_x, l_u
+        (l = -reward), plus a diagonal Gauss-Newton approximation of l_xx/l_uu.
+        """
+        eps = self.fd_eps
+        nx, nu = self.nx, self.nu
+        A = np.zeros((nx, nx))
+        B = np.zeros((nx, nu))
+        l_x = np.zeros(nx)
+        l_u = np.zeros(nu)
+        l_xx = np.zeros(nx)
+        l_uu = np.zeros(nu)
+
+        x0, r0 = self._probe(x, u)
+
+        for i in range(nx):
+            dx = np.zeros(nx); dx[i] = eps
+            xp, rp = self._probe(x + dx, u)
+            xm, rm = self._probe(x - dx, u)
+            A[:, i] = (xp - xm) / (2 * eps)
+            l_x[i] = -(rp - rm) / (2 * eps)
+            l_xx[i] = -(rp - 2 * r0 + rm) / (eps ** 2)
+
+        for j in range(nu):
+            du = np.zeros(nu); du[j] = eps
+            xp, rp = self._probe(x, u + du)
+            xm, rm = self._probe(x, u - du)
+            B[:, j] = (xp - xm) / (2 * eps)
+            l_u[j] = -(rp - rm) / (2 * eps)
+            l_uu[j] = -(rp - 2 * r0 + rm) / (eps ** 2)
+
+        set_state(self.env, x)  # leave env as found
+        return A, B, l_x, l_u, np.diag(np.maximum(l_xx, 0.0)), np.diag(np.maximum(l_uu, 0.0) + self.reg)
+
+    # ------------------------------------------------------------------
+    def _rollout(self, state0: np.ndarray, u_seq: np.ndarray):
+        xs = [state0.copy()]
+        total_cost = 0.0
+        x = state0.copy()
+        for t in range(self.horizon):
+            set_state(self.env, x)
+            u = self._clip_action(u_seq[t])
+            _, r, term, trunc, _ = self.env.step(u)
+            x = np.asarray(get_state(self.env), dtype=float)
+            xs.append(x.copy())
+            total_cost += -float(r)
+            if term or trunc:
+                break
+        set_state(self.env, state0)
+        return xs, total_cost
+
+    def _backward_pass(self, xs: list[np.ndarray], u_seq: np.ndarray):
+        H = len(u_seq)
+        V_x = np.zeros(self.nx)
+        V_xx = np.zeros((self.nx, self.nx))
+        K = [None] * H
+        k = [None] * H
+        for t in reversed(range(H)):
+            A, B, l_x, l_u, l_xx, l_uu = self._linearize(xs[t], u_seq[t])
+            Q_x = l_x + A.T @ V_x
+            Q_u = l_u + B.T @ V_x
+            Q_xx = l_xx + A.T @ V_xx @ A
+            Q_uu = l_uu + B.T @ V_xx @ B
+            Q_ux = B.T @ V_xx @ A
+
+            try:
+                Q_uu_inv = np.linalg.inv(Q_uu)
+            except np.linalg.LinAlgError:
+                Q_uu_inv = np.linalg.pinv(Q_uu)
+
+            K[t] = -Q_uu_inv @ Q_ux
+            k[t] = -Q_uu_inv @ Q_u
+
+            V_x = Q_x + K[t].T @ Q_uu @ k[t] + K[t].T @ Q_u + Q_ux.T @ k[t]
+            V_xx = Q_xx + K[t].T @ Q_uu @ K[t] + K[t].T @ Q_ux + Q_ux.T @ K[t]
+            V_xx = 0.5 * (V_xx + V_xx.T)
+        return K, k
+
+    def predict(self, obs, state=None, deterministic: bool = True):
+        state0 = get_state(self.env)
+        u_seq = self._u_nom.copy()
+        xs, best_cost = self._rollout(state0, u_seq)
+
+        for _ in range(self.iterations):
+            K, k = self._backward_pass(xs[: len(u_seq)], u_seq)
+            improved = False
+            for alpha in self.line_search_steps:
+                if alpha == 0.0:
+                    break  # no-op fallback: keep current nominal
+                u_try = np.zeros_like(u_seq)
+                x = state0.copy()
+                for t in range(len(u_seq)):
+                    u_try[t] = self._clip_action(
+                        u_seq[t] + alpha * k[t] + K[t] @ (x - xs[t])
+                    )
+                    x = self._probe(x, u_try[t])[0]
+                xs_try, cost_try = self._rollout(state0, u_try)
+                if cost_try < best_cost:
+                    u_seq, xs, best_cost = u_try, xs_try, cost_try
+                    improved = True
+                    break
+            if not improved:
+                break
+
+        self._u_nom = u_seq
+        action = self._clip_action(self._u_nom[0])
+        # Receding horizon: shift the plan forward one step.
+        self._u_nom = np.roll(self._u_nom, -1, axis=0)
+        self._u_nom[-1] = 0.0
+        return action, None

tau_ctrl/algorithms/mppi.py

@@ -0,0 +1,265 @@
+"""Sampling-based model-predictive control: MPPI, CEM, and ICEM.
+
+All plan by rolling *candidate action sequences* forward through the env and
+scoring them with the env's **own reward** — so they are fully simulator- and
+reward-agnostic. Rollouts branch from the current state via the env's
+``get_state``/``set_state`` capability (see :mod:`.base`).
+
+Noise: white by default (``noise_beta=0``, the original formulations); set
+``noise_beta`` in (0, 1) for correlated ("colored") noise across the horizon,
+which smooths the resulting action sequences ("Smooth-MPPI"). ``ICEM``
+defaults to colored noise plus elite memory across iterations — see its
+docstring.
+
+Parallelism / GPU: rollouts are embarrassingly parallel. If you pass
+``vector_env_fn`` (a callable returning a Gymnasium ``VectorEnv``), candidates
+are evaluated as a batch — so a GPU-batched simulator env parallelises the
+planner with no code change here. Otherwise a single env is stepped in a loop.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Optional
+
+import numpy as np
+
+from .base import BaseController, ControllerManifest, get_state, register, set_state
+
+
+class _SamplingMPC(BaseController):
+    """Shared machinery for MPPI/CEM: sample sequences, roll out, update."""
+
+    manifest = ControllerManifest(
+        name="sampling_mpc", family="sampling_mpc",
+        requires_branching=True, uses_reward=True, gpu_capable=True,
+    )
+
+    def __init__(
+        self,
+        env: Any,
+        horizon: int = 20,
+        n_samples: int = 100,
+        noise_sigma: float = 1.0,
+        noise_beta: float = 0.0,
+        gamma: float = 1.0,
+        vector_env_fn: Optional[Callable[[int], Any]] = None,
+        **kw: Any,
+    ) -> None:
+        self.horizon = int(horizon)
+        self.n_samples = int(n_samples)
+        self.noise_sigma = float(noise_sigma)
+        self.noise_beta = float(noise_beta)
+        self.gamma = float(gamma)
+        self._vector_env_fn = vector_env_fn
+        super().__init__(env, **kw)
+
+    def _setup(self) -> None:
+        self.nu = self.action_dim
+        self._nominal = np.zeros((self.horizon, self.nu))
+        self._venv = None
+        if self._vector_env_fn is not None:
+            self._venv = self._vector_env_fn(self.n_samples)
+
+    def reset(self) -> None:
+        self._nominal[:] = 0.0
+
+    # --- noise -------------------------------------------------------------
+    def _sample_noise(self, n_samples: int) -> np.ndarray:
+        """Standard-normal noise, shape (n_samples, horizon, nu).
+
+        ``noise_beta = 0`` (default) gives i.i.d. white noise per timestep —
+        vanilla MPPI/CEM. ``noise_beta`` in (0, 1) generates *colored* noise
+        via a first-order AR filter (unit variance preserved), correlating
+        noise across the horizon so sampled action sequences — and the
+        resulting torques — are smooth instead of jittery high-frequency
+        chatter ("Smooth-MPPI", e.g. Vlahov et al.). Higher beta = smoother.
+        """
+        shape = (n_samples, self.horizon, self.nu)
+        if self.noise_beta <= 0:
+            return self.rng.normal(0.0, 1.0, shape)
+        beta = self.noise_beta
+        scale = np.sqrt(1.0 - beta ** 2)
+        noise = np.empty(shape)
+        noise[:, 0, :] = self.rng.normal(0.0, 1.0, (n_samples, self.nu))
+        for t in range(1, self.horizon):
+            noise[:, t, :] = beta * noise[:, t - 1, :] + scale * self.rng.normal(
+                0.0, 1.0, (n_samples, self.nu)
+            )
+        return noise
+
+    # --- rollouts --------------------------------------------------------
+    def _returns(self, state0: Any, seqs: np.ndarray) -> np.ndarray:
+        """Discounted return of each (K, H, nu) action sequence from ``state0``."""
+        if self._venv is not None:
+            return self._returns_vectorized(state0, seqs)
+        return self._returns_serial(state0, seqs)
+
+    def _returns_serial(self, state0: Any, seqs: np.ndarray) -> np.ndarray:
+        K = seqs.shape[0]
+        R = np.zeros(K)
+        for k in range(K):
+            set_state(self.env, state0)
+            disc = 1.0
+            for h in range(self.horizon):
+                _, r, term, trunc, _ = self.env.step(seqs[k, h])
+                R[k] += disc * r
+                disc *= self.gamma
+                if term or trunc:
+                    break
+        set_state(self.env, state0)  # leave the real env as we found it
+        return R
+
+    def _returns_vectorized(self, state0: Any, seqs: np.ndarray) -> np.ndarray:
+        venv = self._venv
+        venv.reset()
+        # Broadcast the current state to every worker, if supported.
+        if hasattr(venv, "set_state"):
+            venv.set_state([state0] * self.n_samples)
+        else:  # per-worker set via the standard call_* API
+            venv.call("set_state", state0)
+        R = np.zeros(self.n_samples)
+        disc = 1.0
+        done = np.zeros(self.n_samples, dtype=bool)
+        for h in range(self.horizon):
+            _, r, term, trunc, _ = venv.step(seqs[:, h])
+            live = ~done
+            R[live] += disc * np.asarray(r)[live]
+            disc *= self.gamma
+            done |= np.asarray(term) | np.asarray(trunc)
+            if done.all():
+                break
+        return R
+
+    # --- update rule (subclass) -----------------------------------------
+    def _update(self, seqs: np.ndarray, returns: np.ndarray) -> None:
+        raise NotImplementedError
+
+    def predict(self, obs, state=None, deterministic: bool = True):
+        state0 = get_state(self.env)
+        noise = self._sample_noise(self.n_samples) * self.noise_sigma
+        seqs = self._clip_seq(self._nominal[None] + noise)
+        returns = self._returns(state0, seqs)
+        set_state(self.env, state0)
+        self._update(seqs, returns)
+
+        action = self._clip_action(self._nominal[0])
+        # Receding horizon: shift the plan forward one step.
+        self._nominal = np.roll(self._nominal, -1, axis=0)
+        self._nominal[-1] = 0.0
+        return action, None
+
+    def _clip_seq(self, seqs: np.ndarray) -> np.ndarray:
+        sp = self.action_space
+        if sp is not None and getattr(sp, "low", None) is not None:
+            return np.clip(seqs, sp.low, sp.high)
+        return seqs
+
+
+@register("mppi")
+class MPPI(_SamplingMPC):
+    """Model-Predictive Path Integral control (softmax-weighted update)."""
+
+    def __init__(self, env: Any, temperature: float = 1.0, **kw: Any) -> None:
+        self.temperature = float(temperature)
+        super().__init__(env, **kw)
+        self.manifest = ControllerManifest(
+            name="mppi", family="sampling_mpc",
+            requires_branching=True, uses_reward=True, gpu_capable=True,
+        )
+
+    def _update(self, seqs: np.ndarray, returns: np.ndarray) -> None:
+        adv = returns - returns.max()  # stabilise the softmax
+        w = np.exp(adv / max(self.temperature, 1e-8))
+        w = w / (w.sum() + 1e-12)
+        self._nominal = np.einsum("k,khu->hu", w, seqs)
+
+
+@register("cem")
+class CEM(_SamplingMPC):
+    """Cross-Entropy Method MPC (refit a Gaussian to the elite sequences)."""
+
+    def __init__(self, env: Any, elite_frac: float = 0.1, iterations: int = 3, **kw: Any) -> None:
+        self.elite_frac = float(elite_frac)
+        self.iterations = int(iterations)
+        super().__init__(env, **kw)
+        self.manifest = ControllerManifest(
+            name="cem", family="sampling_mpc",
+            requires_branching=True, uses_reward=True, gpu_capable=True,
+        )
+
+    def predict(self, obs, state=None, deterministic: bool = True):
+        # CEM refines the distribution over several iterations before acting.
+        state0 = get_state(self.env)
+        mean = self._nominal.copy()
+        sigma = np.full((self.horizon, self.nu), self.noise_sigma)
+        n_elite = max(1, int(self.elite_frac * self.n_samples))
+        for _ in range(self.iterations):
+            noise = self._sample_noise(self.n_samples)
+            seqs = self._clip_seq(mean[None] + sigma[None] * noise)
+            returns = self._returns(state0, seqs)
+            set_state(self.env, state0)
+            elite = seqs[np.argsort(returns)[-n_elite:]]
+            mean = elite.mean(axis=0)
+            sigma = elite.std(axis=0) + 1e-6
+        self._nominal = mean
+        action = self._clip_action(self._nominal[0])
+        self._nominal = np.roll(self._nominal, -1, axis=0)
+        self._nominal[-1] = 0.0
+        return action, None
+
+    def _update(self, seqs, returns):  # unused (predict is overridden)
+        pass
+
+
+@register("icem")
+class ICEM(CEM):
+    """Improved CEM (Pinneri et al. 2020): colored noise + elite memory.
+
+    Two changes over vanilla CEM for the same sample budget:
+
+    1. **Colored noise by default** — correlated across the horizon (see
+       :meth:`_SamplingMPC._sample_noise`), giving smoother action sequences
+       than i.i.d. Gaussian sampling.
+    2. **Elite memory** — the single best sequence found so far *this
+       planning step* is always injected back into the next iteration's
+       population instead of being discarded when the distribution is
+       refit. Vanilla CEM only keeps summary statistics (mean/std) of the
+       elites, so a lucky early sample can be "forgotten"; here it can't be.
+    """
+
+    def __init__(self, env: Any, noise_beta: float = 0.2, **kw: Any) -> None:
+        kw.setdefault("noise_beta", noise_beta)
+        super().__init__(env, **kw)
+        self.manifest = ControllerManifest(
+            name="icem", family="sampling_mpc",
+            requires_branching=True, uses_reward=True, gpu_capable=True,
+        )
+
+    def predict(self, obs, state=None, deterministic: bool = True):
+        state0 = get_state(self.env)
+        mean = self._nominal.copy()
+        sigma = np.full((self.horizon, self.nu), self.noise_sigma)
+        n_elite = max(1, int(self.elite_frac * self.n_samples))
+        best_seq, best_return = None, -np.inf
+
+        for _ in range(self.iterations):
+            noise = self._sample_noise(self.n_samples)
+            seqs = self._clip_seq(mean[None] + sigma[None] * noise)
+            if best_seq is not None:
+                seqs[0] = best_seq  # guarantee the best-so-far is re-evaluated & eligible
+            returns = self._returns(state0, seqs)
+            set_state(self.env, state0)
+
+            order = np.argsort(returns)
+            elite = seqs[order[-n_elite:]]
+            if returns[order[-1]] > best_return:
+                best_return = returns[order[-1]]
+                best_seq = seqs[order[-1]].copy()
+            mean = elite.mean(axis=0)
+            sigma = elite.std(axis=0) + 1e-6
+
+        self._nominal = best_seq if best_seq is not None else mean
+        action = self._clip_action(self._nominal[0])
+        self._nominal = np.roll(self._nominal, -1, axis=0)
+        self._nominal[-1] = 0.0
+        return action, None