tau-ctrl 0.1.0__tar.gz → 0.2.0__tar.gz

{tau_ctrl-0.1.0 → tau_ctrl-0.2.0}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: tau-ctrl
-Version: 0.1.0
-Summary: SB3-like, simulator-agnostic control algorithms over any gymnasium.Env — PID, MPPI/CEM/iCEM, iLQR, CBF, and GPU-native RL (PPO/SAC/TD3) with vectorized on-device training
+Version: 0.2.0
+Summary: Simulator-agnostic control algorithms over any gymnasium.Env — PID, MPPI/CEM/iCEM, iLQR, CBF, and GPU-native RL (PPO/SAC/TD3) with vectorized on-device training
 Author: Tau Intelligence
 License: Apache-2.0
 Project-URL: Homepage, https://tau-intelligence.com/
@@ -40,9 +40,9 @@ Dynamic: license-file
 
 # tau-ctrl
 
-**SB3-like, simulator-agnostic control algorithms — feedback, sampling-based MPC, safety filtering, and GPU-native RL behind one interface.**
+**Simulator-agnostic control algorithms — feedback, sampling-based MPC, safety filtering, and GPU-native RL behind one interface.**
 
-Like Stable-Baselines3, but for the whole controller spectrum, over any `gymnasium.Env`. No simulator dependency: works with whatever hands you an env (e.g. [tau-sim](https://tau-intelligence.com/tau_sim/)). Unlike SB3, the RL methods train **on-device** and scale to vectorized environments for real GPU speedup.
+A unified library for the whole controller spectrum, over any `gymnasium.Env`. No simulator dependency: works with whatever hands you an env (e.g. [tau-sim](https://tau-intelligence.com/tau_sim/)). The RL methods train **on-device** and scale to vectorized environments for real GPU speedup.
 
 ## Installation
 
@@ -57,7 +57,7 @@ pip install tau-ctrl[torch]     # + RL: PPO, SAC, TD3, and vectorized on-device
 from tau_ctrl import make
 
 ctrl = make("mppi", env, horizon=25, n_samples=300)   # or MPPI(env, ...), SAC(env), ...
-action, _ = ctrl.predict(obs)                          # SB3-style
+action, _ = ctrl.predict(obs)                          # standard interface
 ctrl.learn(total_timesteps=100_000)                    # trainable methods (ppo/sac/td3)
 ctrl.save("ctrl.pkl")
 ```
@@ -84,7 +84,7 @@ seeding of torch's RNG).
 
 ## GPU-native RL: vectorized, on-device training
 
-SB3 steps CPU environments and the policy update sits behind per-step Python.
+Standard pipelines step CPU environments sequentially and perform updates in slow Python loops.
 tau-ctrl's SAC/TD3 instead run the replay buffer and the update **on the target
 device**, and — given a batched env — step thousands of environments in parallel
 with no numpy in the hot loop. `Trainer.auto` probes your env and hardware and
@@ -108,8 +108,7 @@ action, _ = model.predict(obs)
 | a single, non-batchable env (PyBullet, classic MuJoCo) + a factory | `SyncTorchVecEnv` | batched update on device |
 | one env you can't replicate (a real robot) | — (single-env path) | only the update |
 
-See [`benchmarks/RESULTS.md`](benchmarks/RESULTS.md) for head-to-head numbers vs
-Stable-Baselines3 and skrl, and [`examples/`](examples/) for runnable scripts
+See [`benchmarks/RESULTS.md`](benchmarks/RESULTS.md) for performance comparison benchmarks, and [`examples/`](examples/) for runnable scripts
 (`quickstart.py`, `vectorized_rl.py`, `adaptive_training.py`).
 
 ## Safety filtering
@@ -151,4 +150,3 @@ Apache 2.0 — see [LICENSE](LICENSE).
 ## Related
 
 - [tau-sim](https://tau-intelligence.com/tau_sim/) — robotics environment builder
-- [Stable-Baselines3](https://stable-baselines3.readthedocs.io/) — RL algorithms (API inspiration)

{tau_ctrl-0.1.0 → tau_ctrl-0.2.0}/README.md

@@ -1,8 +1,8 @@
 # tau-ctrl
 
-**SB3-like, simulator-agnostic control algorithms — feedback, sampling-based MPC, safety filtering, and GPU-native RL behind one interface.**
+**Simulator-agnostic control algorithms — feedback, sampling-based MPC, safety filtering, and GPU-native RL behind one interface.**
 
-Like Stable-Baselines3, but for the whole controller spectrum, over any `gymnasium.Env`. No simulator dependency: works with whatever hands you an env (e.g. [tau-sim](https://tau-intelligence.com/tau_sim/)). Unlike SB3, the RL methods train **on-device** and scale to vectorized environments for real GPU speedup.
+A unified library for the whole controller spectrum, over any `gymnasium.Env`. No simulator dependency: works with whatever hands you an env (e.g. [tau-sim](https://tau-intelligence.com/tau_sim/)). The RL methods train **on-device** and scale to vectorized environments for real GPU speedup.
 
 ## Installation
 
@@ -17,7 +17,7 @@ pip install tau-ctrl[torch]     # + RL: PPO, SAC, TD3, and vectorized on-device
 from tau_ctrl import make
 
 ctrl = make("mppi", env, horizon=25, n_samples=300)   # or MPPI(env, ...), SAC(env), ...
-action, _ = ctrl.predict(obs)                          # SB3-style
+action, _ = ctrl.predict(obs)                          # standard interface
 ctrl.learn(total_timesteps=100_000)                    # trainable methods (ppo/sac/td3)
 ctrl.save("ctrl.pkl")
 ```
@@ -44,7 +44,7 @@ seeding of torch's RNG).
 
 ## GPU-native RL: vectorized, on-device training
 
-SB3 steps CPU environments and the policy update sits behind per-step Python.
+Standard pipelines step CPU environments sequentially and perform updates in slow Python loops.
 tau-ctrl's SAC/TD3 instead run the replay buffer and the update **on the target
 device**, and — given a batched env — step thousands of environments in parallel
 with no numpy in the hot loop. `Trainer.auto` probes your env and hardware and
@@ -68,8 +68,7 @@ action, _ = model.predict(obs)
 | a single, non-batchable env (PyBullet, classic MuJoCo) + a factory | `SyncTorchVecEnv` | batched update on device |
 | one env you can't replicate (a real robot) | — (single-env path) | only the update |
 
-See [`benchmarks/RESULTS.md`](benchmarks/RESULTS.md) for head-to-head numbers vs
-Stable-Baselines3 and skrl, and [`examples/`](examples/) for runnable scripts
+See [`benchmarks/RESULTS.md`](benchmarks/RESULTS.md) for performance comparison benchmarks, and [`examples/`](examples/) for runnable scripts
 (`quickstart.py`, `vectorized_rl.py`, `adaptive_training.py`).
 
 ## Safety filtering
@@ -111,4 +110,3 @@ Apache 2.0 — see [LICENSE](LICENSE).
 ## Related
 
 - [tau-sim](https://tau-intelligence.com/tau_sim/) — robotics environment builder
-- [Stable-Baselines3](https://stable-baselines3.readthedocs.io/) — RL algorithms (API inspiration)

{tau_ctrl-0.1.0 → tau_ctrl-0.2.0}/pyproject.toml

@@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tau-ctrl"
-version = "0.1.0"
-description = "SB3-like, simulator-agnostic control algorithms over any gymnasium.Env — PID, MPPI/CEM/iCEM, iLQR, CBF, and GPU-native RL (PPO/SAC/TD3) with vectorized on-device training"
+version = "0.2.0"
+description = "Simulator-agnostic control algorithms over any gymnasium.Env — PID, MPPI/CEM/iCEM, iLQR, CBF, and GPU-native RL (PPO/SAC/TD3) with vectorized on-device training"
 readme = "README.md"
 license = {text = "Apache-2.0"}
 requires-python = ">=3.9"

{tau_ctrl-0.1.0 → tau_ctrl-0.2.0}/src/tau_ctrl/__init__.py

@@ -1,12 +1,12 @@
 """
-tau-ctrl — an SB3-like, simulator-agnostic controller/algorithm framework.
+tau-ctrl — a simulator-agnostic controller/algorithm framework.
 
-Call algorithms directly off the package, the way you would with SB3:
+Call algorithms directly off the package:
 
     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
+    action, _ = ctrl.predict(obs)                          # returns action, next_state
     ctrl.learn(total_timesteps=100_000)                    # trainable methods (ppo/sac/td3)
 
 Every algorithm shares one interface (``predict``/``learn``/``save``/``load``)
@@ -20,10 +20,15 @@ Model-based methods (MPPI/CEM/ICEM/ILQR/CBF) need the env to be *branchable*
 env. :mod:`tau_ctrl.tuning` adds Bayesian/genetic auto-tuning of any gains.
 """
 
-__version__ = "0.1.0"
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _pkg_version
+
+try:
+    __version__ = _pkg_version("tau-ctrl")
+except PackageNotFoundError:  # not installed (e.g. running from a source checkout)
+    __version__ = "0.0.0+unknown"
 
 from .algorithms import (
-    CBFFilter,
     CEM,
     ICEM,
     ILQR,
@@ -33,7 +38,9 @@ from .algorithms import (
     SAC,
     TD3,
     BaseController,
+    CBFFilter,
     ControllerManifest,
+    MujocoBranchable,
     Trainer,
     available,
     get_state,
@@ -63,6 +70,7 @@ __all__ = [
     "ICEM",
     "ILQR",
     "CBFFilter",
+    "MujocoBranchable",
     "PPO",
     "SAC",
     "TD3",

{tau_ctrl-0.1.0 → tau_ctrl-0.2.0}/src/tau_ctrl/algorithms/__init__.py

@@ -1,12 +1,11 @@
-"""tau-ctrl's SB3-like, simulator-agnostic controller/algorithm framework.
+"""tau-ctrl's simulator-agnostic controller/algorithm framework.
 
-One interface for the whole controller spectrum over a Gymnasium env, called
-the way you'd call SB3:
+One interface for the whole controller spectrum over a Gymnasium env:
 
     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
+    action, _ = ctrl.predict(obs)                          # returns action, next_state
     ctrl.learn(total_timesteps=100_000)                    # for trainable ones (ppo/sac/td3)
 
 Design goals:
@@ -38,6 +37,7 @@ from .base import (
 from .cbf import CBFFilter
 from .ilqr import ILQR
 from .mppi import CEM, ICEM, MPPI
+from .mujoco import MujocoBranchable
 from .pid import PID
 from .ppo import PPO
 from .sac import SAC
@@ -72,6 +72,7 @@ __all__ = [
     "ICEM",
     "ILQR",
     "CBFFilter",
+    "MujocoBranchable",
     "PPO",
     "SAC",
     "TD3",

{tau_ctrl-0.1.0 → tau_ctrl-0.2.0}/src/tau_ctrl/algorithms/base.py

@@ -1,9 +1,9 @@
-"""Core abstractions for tau-ctrl's algorithm framework — an SB3-like,
-simulator-agnostic controller framework over a Gymnasium environment.
+"""Core abstractions for tau-ctrl's algorithm framework — a 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
+same :class:`BaseController` interface — ``predict(obs) -> action``, 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
@@ -37,16 +37,78 @@ def register(name: str):
 
 
 def make(name: str, env: Any, **kwargs: Any) -> "BaseController":
-    """Construct a controller by name, SB3-style: ``make("mppi", env, ...)``."""
+    """Construct a controller by name: ``make("mppi", env, ...)``."""
     if name not in _REGISTRY:
         raise KeyError(f"unknown controller {name!r}; available: {sorted(_REGISTRY)}")
+    _reject_raw_vector_env(env)
     return _REGISTRY[name](env, **kwargs)
 
 
+def _reject_raw_vector_env(env: Any) -> None:
+    """Guard against handing a raw ``gymnasium.vector.VectorEnv`` to a controller.
+
+    Its *batched* spaces would size the networks wrong (``obs_dim`` would come
+    out as ``num_envs * obs_dim``) and its ``reset``/``step`` signature doesn't
+    match the on-device vectorized trainer, so ``learn()`` would crash. Route
+    vectorized envs through :meth:`Trainer.auto`, which wraps them
+    (``GymVectorAdapter``) and picks a convergence-safe update ratio. A native
+    ``tau_ctrl`` ``TorchVecEnv`` (incl. the adapters) is fine and passes through.
+    """
+    try:
+        import gymnasium  # noqa: PLC0415
+
+        is_gym_vec = isinstance(env, gymnasium.vector.VectorEnv)
+    except Exception:
+        is_gym_vec = False
+    if is_gym_vec:
+        raise TypeError(
+            "make() received a gymnasium.vector.VectorEnv, which cannot be used "
+            "directly (its batched spaces mis-size the networks). Train vectorized "
+            "envs via Trainer.auto, which adapts them to the on-device loop:\n"
+            "    from tau_ctrl import Trainer\n"
+            "    model = Trainer.auto('sac', env=vec_env, total_timesteps=...)"
+        )
+
+
 def available() -> list[str]:
     return sorted(_REGISTRY)
 
 
+def require_torch():
+    """Import and return torch, or raise a clear, actionable error.
+
+    The RL methods (PPO/SAC/TD3) and ``Trainer.auto`` need PyTorch, which ships
+    only via the optional extra. Without this, a user hits a bare
+    ``ModuleNotFoundError: No module named 'torch'`` with no hint of the fix.
+    """
+    try:
+        import torch  # noqa: PLC0415
+    except ImportError as e:
+        raise ImportError(
+            "PyTorch is required for the RL methods (PPO/SAC/TD3) and Trainer.auto, "
+            "but it is not installed. Install it with:\n"
+            "    pip install tau-ctrl[torch]"
+        ) from e
+    return torch
+
+
+def require_box_action_space(action_space: Any, algo: str) -> None:
+    """Reject non-continuous action spaces up front with a clear message.
+
+    PID and the RL/MPC controllers all assume a continuous ``Box`` action
+    space; handing them a ``Discrete`` env otherwise fails deep inside with an
+    opaque ``AttributeError: 'Discrete' object has no attribute 'low'``.
+    """
+    from gymnasium import spaces  # noqa: PLC0415
+
+    if not isinstance(action_space, spaces.Box):
+        raise TypeError(
+            f"{algo} requires a continuous Box action space, got "
+            f"{type(action_space).__name__}. These controllers are for continuous "
+            "control; discrete-action envs are not supported."
+        )
+
+
 # ---------------------------------------------------------------------------
 # Manifest
 # ---------------------------------------------------------------------------
@@ -73,23 +135,31 @@ class BranchingNotSupported(RuntimeError):
 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``.
+    Resolution order (no simulator assumptions): an explicit ``get_state``
+    anywhere on the wrapper stack — the env itself or a branchable wrapper such
+    as :class:`~tau_ctrl.algorithms.mujoco.MujocoBranchable` layered over a sim
+    that doesn't expose state natively — else classic-control ``.state``.
+    Simulators that want model-based control just expose ``get_state``/``set_state``.
     """
     u = getattr(env, "unwrapped", env)
+    if hasattr(env, "get_state"):
+        return env.get_state()
     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."
+        "(MPPI/CEM/CBF) need a branchable env. For MuJoCo envs, wrap with "
+        "tau_ctrl.MujocoBranchable(env)."
     )
 
 
 def set_state(env: Any, state: Any) -> None:
     u = getattr(env, "unwrapped", env)
+    if hasattr(env, "set_state"):
+        env.set_state(state)
+        return
     if hasattr(u, "set_state"):
         u.set_state(state)
         return
@@ -131,7 +201,7 @@ def resolve_device(pref: str = "auto") -> str:
 # ---------------------------------------------------------------------------
 
 class BaseController(ABC):
-    """SB3-like controller: ``predict`` always; ``learn`` when trainable."""
+    """Base controller: ``predict`` always; ``learn`` when trainable."""
 
     name: str = "base"
     manifest: ControllerManifest
@@ -164,7 +234,7 @@ class BaseController(ABC):
     def predict(
         self, obs: Any, state: Any = None, deterministic: bool = True
     ) -> tuple[np.ndarray, Any]:
-        """Return ``(action, next_internal_state)`` — SB3 signature."""
+        """Return ``(action, next_internal_state)``."""
 
     def learn(self, total_timesteps: int = 0, **kwargs: Any) -> "BaseController":
         """Train the controller. Analytic controllers are no-ops."""

{tau_ctrl-0.1.0 → tau_ctrl-0.2.0}/src/tau_ctrl/algorithms/ilqr.py

@@ -56,6 +56,11 @@ class ILQR(BaseController):
         state0 = get_state(self.env)
         self.nx = int(np.asarray(state0).size)
         self._u_nom = np.zeros((self.horizon, self.nu))
+        # Probe/roll out against the *unwrapped* env (see MPPI): stepping the
+        # wrapped env would trip TimeLimit truncation mid-rollout (producing a
+        # short trajectory and an IndexError in the backward pass) and corrupt
+        # the live episode's wrapper state.
+        self._rollout_env = getattr(self.env, "unwrapped", self.env)
 
     def reset(self) -> None:
         self._u_nom[:] = 0.0
@@ -67,7 +72,7 @@ class ILQR(BaseController):
         """One step from (x, u): returns (next_state, reward)."""
         set_state(self.env, x)
         u = self._clip_action(u)
-        _, r, _, _, _ = self.env.step(u)
+        _, r, _, _, _ = self._rollout_env.step(u)
         return np.asarray(get_state(self.env), dtype=float), float(r)
 
     def _linearize(self, x: np.ndarray, u: np.ndarray):
@@ -112,7 +117,7 @@ class ILQR(BaseController):
         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)
+            _, r, term, trunc, _ = self._rollout_env.step(u)
             x = np.asarray(get_state(self.env), dtype=float)
             xs.append(x.copy())
             total_cost += -float(r)
@@ -154,14 +159,21 @@ class ILQR(BaseController):
         xs, best_cost = self._rollout(state0, u_seq)
 
         for _ in range(self.iterations):
-            K, k = self._backward_pass(xs[: len(u_seq)], u_seq)
+            # A rollout can terminate before the full horizon (a genuinely
+            # terminating env, e.g. the ant falling), leaving fewer states than
+            # controls. Only optimize over the segment we actually have a
+            # trajectory for; controls past it stay at their nominal value.
+            H = min(len(u_seq), len(xs) - 1)
+            if H <= 0:
+                break
+            K, k = self._backward_pass(xs[:H], u_seq[:H])
             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)
+                u_try = u_seq.copy()
                 x = state0.copy()
-                for t in range(len(u_seq)):
+                for t in range(H):
                     u_try[t] = self._clip_action(
                         u_seq[t] + alpha * k[t] + K[t] @ (x - xs[t])
                     )

{tau_ctrl-0.1.0 → tau_ctrl-0.2.0}/src/tau_ctrl/algorithms/mppi.py

@@ -56,6 +56,12 @@ class _SamplingMPC(BaseController):
     def _setup(self) -> None:
         self.nu = self.action_dim
         self._nominal = np.zeros((self.horizon, self.nu))
+        # Roll candidate sequences out against the *unwrapped* env. get_state /
+        # set_state already operate on env.unwrapped, so stepping the wrapped env
+        # here would leak into stateful wrappers (TimeLimit._elapsed_steps,
+        # OrderEnforcing) and corrupt/terminate the live episode — while the
+        # dynamics and reward are identical either way.
+        self._rollout_env = getattr(self.env, "unwrapped", self.env)
         self._venv = None
         if self._vector_env_fn is not None:
             self._venv = self._vector_env_fn(self.n_samples)
@@ -101,7 +107,7 @@ class _SamplingMPC(BaseController):
             set_state(self.env, state0)
             disc = 1.0
             for h in range(self.horizon):
-                _, r, term, trunc, _ = self.env.step(seqs[k, h])
+                _, r, term, trunc, _ = self._rollout_env.step(seqs[k, h])
                 R[k] += disc * r
                 disc *= self.gamma
                 if term or trunc:

tau_ctrl-0.2.0/src/tau_ctrl/algorithms/mujoco.py

@@ -0,0 +1,72 @@
+"""Make a MuJoCo Gymnasium env *branchable* for the model-based controllers.
+
+The sampling/gradient planners (MPPI/CEM/ICEM/ILQR) and the CBF safety filter
+plan by rolling candidate action sequences forward from the current state, which
+needs the env to expose ``get_state``/``set_state`` (see
+:mod:`tau_ctrl.algorithms.base`). Classic-control envs provide this via
+``.state``; MuJoCo envs don't. :class:`MujocoBranchable` adds it by
+snapshotting/restoring MuJoCo's *full physics state* (time, ``qpos``, ``qvel``,
+actuator activation, plugin state) through ``mj_getState``/``mj_setState`` — the
+complete, deterministic state, not just ``qpos``/``qvel``::
+
+    import gymnasium as gym
+    from tau_ctrl import make, MujocoBranchable
+
+    env = MujocoBranchable(gym.make("HalfCheetah-v5"))
+    ctrl = make("mppi", env, horizon=20, n_samples=200)   # now branchable
+    action, _ = ctrl.predict(obs)
+
+Scope
+-----
+Works on **any** Gymnasium MuJoCo env — the built-in ones (HalfCheetah, Ant,
+Hopper, …) and your own ``MujocoEnv`` subclasses — whose complete state is the
+physics state. It does **not** magically make a *non*-MuJoCo env branchable:
+branching needs the full state, which an arbitrary env (hidden RNG, external
+process) can't expose. If a *custom* MuJoCo env also carries Python-side state
+that affects dynamics or reward (a goal, a step counter, an action history),
+subclass this and extend ``get_state``/``set_state`` to include it.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+
+
+class MujocoBranchable(gym.Wrapper):
+    """Wrap a MuJoCo Gymnasium env so model-based controllers can branch it.
+
+    Adds ``get_state``/``set_state`` backed by MuJoCo's full physics state.
+    Transparent otherwise — ``reset``/``step``/spaces pass straight through.
+    """
+
+    def __init__(self, env: gym.Env) -> None:
+        super().__init__(env)
+        raw = env.unwrapped
+        if not (hasattr(raw, "model") and hasattr(raw, "data")):
+            raise TypeError(
+                "MujocoBranchable expects a MuJoCo Gymnasium env (env.unwrapped "
+                f"with .model/.data), got {type(raw).__name__}. Non-MuJoCo envs "
+                "must provide get_state/set_state themselves to be branchable."
+            )
+        import mujoco  # noqa: PLC0415
+
+        self._mujoco = mujoco
+        self._spec = mujoco.mjtState.mjSTATE_FULLPHYSICS
+        self._size = int(mujoco.mj_stateSize(raw.model, self._spec))
+
+    def get_state(self) -> np.ndarray:
+        raw = self.env.unwrapped
+        arr = np.zeros(self._size, dtype=np.float64)
+        self._mujoco.mj_getState(raw.model, raw.data, arr, self._spec)
+        return arr
+
+    def set_state(self, state: Any) -> None:
+        raw = self.env.unwrapped
+        arr = np.ascontiguousarray(state, dtype=np.float64)
+        self._mujoco.mj_setState(raw.model, raw.data, arr, self._spec)
+        # Recompute derived quantities (contacts, sensors) so the restored state
+        # is fully consistent before the next step is planned/taken.
+        self._mujoco.mj_forward(raw.model, raw.data)

{tau_ctrl-0.1.0 → tau_ctrl-0.2.0}/src/tau_ctrl/algorithms/ppo.py

@@ -1,8 +1,8 @@
 """A compact, self-contained PPO (continuous actions) in PyTorch.
 
-SB3-like surface (``learn`` / ``predict`` / ``save`` / ``load``) but our own
-implementation so tau-ctrl doesn't depend on stable-baselines3. Runs on GPU
-automatically when torch reports CUDA available (``device="auto"``), else CPU.
+Standard surface (``learn`` / ``predict`` / ``save`` / ``load``) implementation.
+Runs on GPU automatically when torch reports CUDA available (``device="auto"``),
+else CPU.
 
 Simulator-agnostic: talks only to the Gymnasium env's ``reset``/``step``.
 Kept intentionally small (single-env on-policy rollouts, GAE, clipped
@@ -16,12 +16,17 @@ from typing import Any
 
 import numpy as np
 
-from .base import BaseController, ControllerManifest, register
+from .base import (
+    BaseController,
+    ControllerManifest,
+    register,
+    require_box_action_space,
+    require_torch,
+)
 
 
 def _torch():
-    import torch  # noqa: PLC0415
-    return torch
+    return require_torch()
 
 
 @register("ppo")
@@ -52,6 +57,7 @@ class PPO(BaseController):
         super().__init__(env, **kw)
 
     def _setup(self) -> None:
+        require_box_action_space(self.action_space, "PPO")
         torch = _torch()
         import torch.nn as nn
 

{tau_ctrl-0.1.0 → tau_ctrl-0.2.0}/src/tau_ctrl/algorithms/sac.py

@@ -17,7 +17,13 @@ from typing import Any, Optional
 
 import numpy as np
 
-from .base import BaseController, ControllerManifest, register
+from .base import (
+    BaseController,
+    ControllerManifest,
+    register,
+    require_box_action_space,
+    require_torch,
+)
 from .off_policy import (
     ReplayBuffer,
     build_mlp,
@@ -29,8 +35,7 @@ from .off_policy import (
 
 
 def _torch():
-    import torch  # noqa: PLC0415
-    return torch
+    return require_torch()
 
 
 @register("sac")
@@ -57,6 +62,7 @@ class SAC(BaseController):
         super().__init__(env, **kw)
 
     def _setup(self) -> None:
+        require_box_action_space(self.action_space, "SAC")
         torch = _torch()
         import torch.nn as nn
 

{tau_ctrl-0.1.0 → tau_ctrl-0.2.0}/src/tau_ctrl/algorithms/strategy.py

@@ -2,15 +2,15 @@
 execution strategy, dispatch. One ``Trainer.auto(...)`` call adapts to whatever
 you hand it, so tau-ctrl generalizes past any single env.
 
-SB3 has one execution model (CPU vector env, synchronous collect-then-update)
-regardless of what env you give it. The whole speed argument for tau-ctrl is
-that wall-clock-to-reward is bounded by the *slower* of collection and update,
-and which one dominates depends on the env — so a general framework must detect
-the regime and route to the matching engine:
+Standard frameworks typically support a single execution model (e.g., CPU vector envs,
+synchronous collect-then-update) regardless of the target environment. The whole
+speed argument for tau-ctrl is that wall-clock-to-reward is bounded by the *slower* of
+collection and update, and which one dominates depends on the env — so a general
+framework must detect the regime and route to the matching engine:
 
     ON_DEVICE_VECTORIZED  env already batched + on GPU (MJX/Brax via dlpack,
-                          Isaac Gym, or a native TorchVecEnv) → 10-100x, the
-                          real win; SB3 can't do this at all.
+                          Isaac Gym, or a native TorchVecEnv) → 10-100x speedup
+                          by running everything on-device.
     GYM_VECTOR_ADAPTED    a gymnasium.vector.VectorEnv (numpy/CPU) → batched
                           update on device, collection stays CPU.
     SYNC_VEC              a single non-batchable env replicated N times, stepped
@@ -192,7 +192,7 @@ def select_strategy(
         why = (
             f"Env is a native on-device TorchVecEnv ({ec.backend}, on_device={ec.on_device}). "
             "Best case: env stepping AND the gradient update run batched on the target "
-            "device — the regime where tau-ctrl beats SB3 10-100x (SB3's VecEnv can't). "
+            "device — the regime that maximizes GPU throughput by avoiding CPU-GPU synchronization. "
             f"Scaling updates_per_step→{ups} so {n} fresh transitions/step actually train."
         )
         return Plan(Strategy.ON_DEVICE_VECTORIZED, n, dev, ups, why)
@@ -225,7 +225,7 @@ def select_strategy(
 
     why = (
         "Single env, not replicable (e.g. a real robot). Physics can't be parallelised — "
-        "only the gradient update is GPU-accelerated. Honest regime: beat SB3 here on "
+        "only the gradient update is GPU-accelerated. Honest regime: focus on "
         "sample efficiency (fewer env steps to reward), not raw throughput. Use SAC/TD3."
     )
     return Plan(Strategy.SINGLE_ENV, 1, dev, 1, why)

{tau_ctrl-0.1.0 → tau_ctrl-0.2.0}/src/tau_ctrl/algorithms/td3.py

@@ -13,7 +13,13 @@ from typing import Any
 
 import numpy as np
 
-from .base import BaseController, ControllerManifest, register
+from .base import (
+    BaseController,
+    ControllerManifest,
+    register,
+    require_box_action_space,
+    require_torch,
+)
 from .off_policy import (
     ReplayBuffer,
     build_mlp,
@@ -25,8 +31,7 @@ from .off_policy import (
 
 
 def _torch():
-    import torch  # noqa: PLC0415
-    return torch
+    return require_torch()
 
 
 @register("td3")
@@ -57,6 +62,7 @@ class TD3(BaseController):
         super().__init__(env, **kw)
 
     def _setup(self) -> None:
+        require_box_action_space(self.action_space, "TD3")
         torch = _torch()
         import torch.nn as nn
 

{tau_ctrl-0.1.0 → tau_ctrl-0.2.0}/src/tau_ctrl/algorithms/vec_env.py

@@ -1,9 +1,9 @@
 """GPU-vectorized environments — the piece that unlocks real GPU speedup.
 
-SB3 (and tau-ctrl's single-env ``learn()``) step one CPU environment at a
-time: one transition, one tiny update, ~100 sequential CUDA kernel launches
-per step.  At that granularity kernel-launch latency dwarfs the actual math,
-so "putting the network on the GPU" caps out at ~1.3x.
+Single-environment training loops (and tau-ctrl's single-env ``learn()``) step
+one CPU environment at a time: one transition, one tiny update, ~100 sequential
+CUDA kernel launches per step.  At that granularity kernel-launch latency dwarfs
+the actual math, so "putting the network on the GPU" caps out at ~1.3x.
 
 A :class:`TorchVecEnv` instead advances ``num_envs`` environments *in
 parallel*, entirely as batched tensor ops on the target device.  Nothing

{tau_ctrl-0.1.0 → tau_ctrl-0.2.0}/src/tau_ctrl.egg-info/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: tau-ctrl
-Version: 0.1.0
-Summary: SB3-like, simulator-agnostic control algorithms over any gymnasium.Env — PID, MPPI/CEM/iCEM, iLQR, CBF, and GPU-native RL (PPO/SAC/TD3) with vectorized on-device training
+Version: 0.2.0
+Summary: Simulator-agnostic control algorithms over any gymnasium.Env — PID, MPPI/CEM/iCEM, iLQR, CBF, and GPU-native RL (PPO/SAC/TD3) with vectorized on-device training
 Author: Tau Intelligence
 License: Apache-2.0
 Project-URL: Homepage, https://tau-intelligence.com/
@@ -40,9 +40,9 @@ Dynamic: license-file
 
 # tau-ctrl
 
-**SB3-like, simulator-agnostic control algorithms — feedback, sampling-based MPC, safety filtering, and GPU-native RL behind one interface.**
+**Simulator-agnostic control algorithms — feedback, sampling-based MPC, safety filtering, and GPU-native RL behind one interface.**
 
-Like Stable-Baselines3, but for the whole controller spectrum, over any `gymnasium.Env`. No simulator dependency: works with whatever hands you an env (e.g. [tau-sim](https://tau-intelligence.com/tau_sim/)). Unlike SB3, the RL methods train **on-device** and scale to vectorized environments for real GPU speedup.
+A unified library for the whole controller spectrum, over any `gymnasium.Env`. No simulator dependency: works with whatever hands you an env (e.g. [tau-sim](https://tau-intelligence.com/tau_sim/)). The RL methods train **on-device** and scale to vectorized environments for real GPU speedup.
 
 ## Installation
 
@@ -57,7 +57,7 @@ pip install tau-ctrl[torch]     # + RL: PPO, SAC, TD3, and vectorized on-device
 from tau_ctrl import make
 
 ctrl = make("mppi", env, horizon=25, n_samples=300)   # or MPPI(env, ...), SAC(env), ...
-action, _ = ctrl.predict(obs)                          # SB3-style
+action, _ = ctrl.predict(obs)                          # standard interface
 ctrl.learn(total_timesteps=100_000)                    # trainable methods (ppo/sac/td3)
 ctrl.save("ctrl.pkl")
 ```
@@ -84,7 +84,7 @@ seeding of torch's RNG).
 
 ## GPU-native RL: vectorized, on-device training
 
-SB3 steps CPU environments and the policy update sits behind per-step Python.
+Standard pipelines step CPU environments sequentially and perform updates in slow Python loops.
 tau-ctrl's SAC/TD3 instead run the replay buffer and the update **on the target
 device**, and — given a batched env — step thousands of environments in parallel
 with no numpy in the hot loop. `Trainer.auto` probes your env and hardware and
@@ -108,8 +108,7 @@ action, _ = model.predict(obs)
 | a single, non-batchable env (PyBullet, classic MuJoCo) + a factory | `SyncTorchVecEnv` | batched update on device |
 | one env you can't replicate (a real robot) | — (single-env path) | only the update |
 
-See [`benchmarks/RESULTS.md`](benchmarks/RESULTS.md) for head-to-head numbers vs
-Stable-Baselines3 and skrl, and [`examples/`](examples/) for runnable scripts
+See [`benchmarks/RESULTS.md`](benchmarks/RESULTS.md) for performance comparison benchmarks, and [`examples/`](examples/) for runnable scripts
 (`quickstart.py`, `vectorized_rl.py`, `adaptive_training.py`).
 
 ## Safety filtering
@@ -151,4 +150,3 @@ Apache 2.0 — see [LICENSE](LICENSE).
 ## Related
 
 - [tau-sim](https://tau-intelligence.com/tau_sim/) — robotics environment builder
-- [Stable-Baselines3](https://stable-baselines3.readthedocs.io/) — RL algorithms (API inspiration)

{tau_ctrl-0.1.0 → tau_ctrl-0.2.0}/src/tau_ctrl.egg-info/SOURCES.txt

@@ -14,6 +14,7 @@ src/tau_ctrl/algorithms/base.py
 src/tau_ctrl/algorithms/cbf.py
 src/tau_ctrl/algorithms/ilqr.py
 src/tau_ctrl/algorithms/mppi.py
+src/tau_ctrl/algorithms/mujoco.py
 src/tau_ctrl/algorithms/off_policy.py
 src/tau_ctrl/algorithms/pid.py
 src/tau_ctrl/algorithms/ppo.py