diff --git a/pyproject.toml b/pyproject.toml index f1a7090..f075644 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ authors = [ keywords = [ "robots", "agents", "strands", "sdk", "mujoco", "simulation", "lerobot", "reinforcement-learning", - "gr00t", "isaac-sim", "vla", + "gr00t", "isaac-sim", "vla", "newton", "warp", "robot-learning", "imitation-learning", "teleoperation", "manipulation", "strands-agents", ] @@ -51,10 +51,15 @@ lerobot = [ sim = [ "robot_descriptions>=1.11.0,<2.0.0", ] +newton = [ + "warp-lang>=1.1.0,<2.0.0", + "newton-sim>=0.4.0,<1.0.0", +] all = [ "strands-robots[groot-service]", "strands-robots[lerobot]", "strands-robots[sim]", + # "strands-robots[newton]", # excluded: newton-sim not yet on PyPI ] dev = [ "pytest>=6.0,<9.0.0", @@ -128,7 +133,7 @@ ignore_missing_imports = false # Third-party libs without type stubs [[tool.mypy.overrides]] -module = ["lerobot.*", "gr00t.*", "draccus.*", "msgpack.*", "zmq.*", "huggingface_hub.*", "serial.*", "psutil.*", "torch.*", "torchvision.*", "transformers.*", "einops.*", "robot_descriptions.*"] +module = ["lerobot.*", "gr00t.*", "draccus.*", "msgpack.*", "zmq.*", "huggingface_hub.*", "serial.*", "psutil.*", "torch.*", "torchvision.*", "transformers.*", "einops.*", "robot_descriptions.*", "warp.*", "newton.*"] ignore_missing_imports = true # @tool decorator injects runtime signatures mypy cannot check diff --git a/strands_robots/simulation/factory.py b/strands_robots/simulation/factory.py index e7b0a5b..203161b 100644 --- a/strands_robots/simulation/factory.py +++ b/strands_robots/simulation/factory.py @@ -13,9 +13,8 @@ # Explicit backend sim = create_simulation("mujoco", timestep=0.001) - # Future backends - sim = create_simulation("isaac", gpu_id=0) - sim = create_simulation("newton") + # GPU-native backend + sim = create_simulation("newton", num_envs=4096) # Custom backend (runtime-registered) from strands_robots.simulation.factory import register_backend @@ -43,15 +42,20 @@ "strands_robots.simulation.mujoco.simulation", "Simulation", ), + "newton": ( + "strands_robots.simulation.newton.simulation", + "NewtonSimulation", + ), # Future: # "isaac": ("strands_robots.simulation.isaac.simulation", "IsaacSimulation"), - # "newton": ("strands_robots.simulation.newton.simulation", "NewtonSimulation"), } _BUILTIN_ALIASES: dict[str, str] = { "mj": "mujoco", "mjc": "mujoco", "mjx": "mujoco", + "warp": "newton", + "wp": "newton", # "isaac_sim": "isaac", # "isaacsim": "isaac", # "nvidia": "isaac", @@ -139,7 +143,7 @@ def list_backends() -> list[str]: Example:: >>> list_backends() - ['mj', 'mjc', 'mjx', 'mujoco'] + ['mj', 'mjc', 'mjx', 'mujoco', 'newton', 'warp', 'wp'] """ names: set[str] = set() names.update(_BUILTIN_BACKENDS.keys()) @@ -201,9 +205,11 @@ def create_simulation( Args: backend: Backend name or alias. Defaults to ``"mujoco"``. - Built-in: ``"mujoco"`` (aliases: ``"mj"``, ``"mjc"``, ``"mjx"``). + Built-in: ``"mujoco"`` (aliases: ``"mj"``, ``"mjc"``, ``"mjx"``), + ``"newton"`` (aliases: ``"warp"``, ``"wp"``). **kwargs: Backend-specific keyword arguments passed to the - constructor (e.g., ``tool_name``, ``timestep``). + constructor (e.g., ``tool_name``, ``timestep``, + ``num_envs``, ``solver``). Returns: A ``SimEngine`` instance ready for ``create_world()``. @@ -211,7 +217,7 @@ def create_simulation( Raises: ValueError: If the backend name is not recognized. ImportError: If the backend's dependencies are missing - (e.g., ``pip install mujoco``). + (e.g., ``pip install mujoco`` or ``pip install 'strands-robots[newton]'``). Examples:: @@ -220,8 +226,11 @@ def create_simulation( sim.create_world() sim.add_robot("so100") + # Newton GPU backend + sim = create_simulation("newton", num_envs=4096) + # With alias - sim = create_simulation("mj") + sim = create_simulation("warp") # Pass kwargs to backend constructor sim = create_simulation("mujoco", tool_name="my_sim") diff --git a/strands_robots/simulation/newton/__init__.py b/strands_robots/simulation/newton/__init__.py new file mode 100644 index 0000000..8d27d46 --- /dev/null +++ b/strands_robots/simulation/newton/__init__.py @@ -0,0 +1,47 @@ +"""Newton simulation backend — GPU-native physics via NVIDIA Warp + Newton. + +Lazy-loading module: importing this package does NOT trigger ``import warp`` +or ``import newton``. Heavy dependencies are loaded only when +``NewtonSimulation`` is instantiated or its methods are called. + +Usage:: + + from strands_robots.simulation.newton import NewtonSimulation, NewtonConfig + + config = NewtonConfig(num_envs=4096, solver="mujoco", device="cuda:0") + sim = NewtonSimulation(config=config) + +Or via the factory:: + + from strands_robots.simulation import create_simulation + sim = create_simulation("newton", num_envs=4096) +""" + +from __future__ import annotations + +from typing import Any + +# Light import — dataclass, no heavy deps +from strands_robots.simulation.newton.config import NewtonConfig + +# Lazy-loaded heavy imports +_LAZY_IMPORTS: dict[str, tuple[str, str]] = { + "NewtonSimulation": ( + "strands_robots.simulation.newton.simulation", + "NewtonSimulation", + ), +} + +__all__ = ["NewtonConfig", "NewtonSimulation"] + + +def __getattr__(name: str) -> Any: + if name in _LAZY_IMPORTS: + import importlib + + module_path, attr_name = _LAZY_IMPORTS[name] + module = importlib.import_module(module_path) + value = getattr(module, attr_name) + globals()[name] = value + return value + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/strands_robots/simulation/newton/config.py b/strands_robots/simulation/newton/config.py new file mode 100644 index 0000000..bfe662e --- /dev/null +++ b/strands_robots/simulation/newton/config.py @@ -0,0 +1,105 @@ +"""Newton backend configuration. + +Dataclass-only module — no heavy dependencies (no warp, no newton). +Safe to import at module level without triggering GPU initialization. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +# Solver name → Newton solver class name. +# Kept here (not solvers.py) so config validation works without importing warp. +SOLVER_MAP: dict[str, str] = { + "mujoco": "SolverMuJoCo", + "featherstone": "SolverFeatherstone", + "semi_implicit": "SolverSemiImplicit", + "xpbd": "SolverXPBD", + "vbd": "SolverVBD", + "style3d": "SolverStyle3D", + "implicit_mpm": "SolverImplicitMPM", +} + +RENDER_BACKENDS: frozenset[str] = frozenset({"opengl", "rerun", "viser", "null", "none"}) + +BROAD_PHASE_OPTIONS: frozenset[str] = frozenset({"sap", "bvh", "none"}) + + +@dataclass +class NewtonConfig: + """Configuration for the Newton GPU simulation backend. + + Parameters + ---------- + num_envs : int + Number of parallel environments. GPU backends support 4096+ + on a single device. + device : str + Warp device string (``"cuda:0"``, ``"cpu"``). + solver : str + Solver backend name. Must be a key in ``SOLVER_MAP``. + ``"mujoco"`` is the default — fastest for rigid-body articulated + systems. ``"xpbd"`` / ``"semi_implicit"`` are alternatives. + ``"vbd"`` / ``"style3d"`` are cloth-only. ``"implicit_mpm"`` is + for granular/fluid simulation. + physics_dt : float + Physics timestep in seconds. + substeps : int + Number of substeps per ``step()`` call. + render_backend : str + Rendering backend (``"opengl"``, ``"rerun"``, ``"viser"``, + ``"null"``/``"none"``). + enable_cuda_graph : bool + Capture the simulation loop in a CUDA graph for minimal + Python overhead. Requires static graph (no dynamic shapes). + enable_differentiable : bool + Enable gradient tracking for differentiable simulation + via ``wp.Tape``. + broad_phase : str + Broad-phase collision algorithm (``"sap"``, ``"bvh"``, ``"none"``). + soft_contact_margin : float + Soft-contact margin distance. + soft_contact_ke : float + Contact stiffness. + soft_contact_kd : float + Contact damping. + soft_contact_mu : float + Contact friction coefficient. + soft_contact_restitution : float + Contact restitution coefficient. + + Raises + ------ + ValueError + If ``solver``, ``render_backend``, or ``broad_phase`` is invalid, + or if ``physics_dt <= 0`` or ``num_envs < 1``. + """ + + num_envs: int = 1 + device: str = "cuda:0" + solver: str = "mujoco" + physics_dt: float = 0.005 + substeps: int = 1 + render_backend: str = "none" + enable_cuda_graph: bool = False + enable_differentiable: bool = False + broad_phase: str = "sap" + soft_contact_margin: float = 0.5 + soft_contact_ke: float = 10000.0 + soft_contact_kd: float = 10.0 + soft_contact_mu: float = 0.5 + soft_contact_restitution: float = 0.0 + + def __post_init__(self) -> None: + if self.solver not in SOLVER_MAP: + raise ValueError(f"Unknown solver {self.solver!r}. Valid options: {sorted(SOLVER_MAP.keys())}") + if self.render_backend not in RENDER_BACKENDS: + raise ValueError( + f"Unknown render_backend {self.render_backend!r}. Valid options: {sorted(RENDER_BACKENDS)}" + ) + if self.broad_phase not in BROAD_PHASE_OPTIONS: + raise ValueError(f"Unknown broad_phase {self.broad_phase!r}. Valid options: {sorted(BROAD_PHASE_OPTIONS)}") + if self.physics_dt <= 0: + raise ValueError(f"physics_dt must be positive, got {self.physics_dt}") + if self.num_envs < 1: + raise ValueError(f"num_envs must be >= 1, got {self.num_envs}") diff --git a/strands_robots/simulation/newton/simulation.py b/strands_robots/simulation/newton/simulation.py new file mode 100644 index 0000000..1057bd1 --- /dev/null +++ b/strands_robots/simulation/newton/simulation.py @@ -0,0 +1,412 @@ +"""Newton GPU simulation backend — SimEngine implementation. + +This module contains the ``NewtonSimulation`` class, which implements +the ``SimEngine`` ABC for the NVIDIA Warp + Newton physics engine. + +Heavy dependencies (``warp``, ``newton``) are imported lazily on +first use — importing this module alone is lightweight. + +.. note:: + + This is the **skeleton** PR. Methods raise ``NotImplementedError`` + where actual Newton API calls will go. Subsequent PRs will fill in + the implementations one category at a time: + + - PR 2: world lifecycle + robot loading + - PR 3: step / action / observation + - PR 4: object management + rendering + - PR 5: replicate, soft bodies + - PR 6: diffsim, IK, sensors +""" + +from __future__ import annotations + +import logging +from typing import Any + +from strands_robots.simulation.base import SimEngine +from strands_robots.simulation.newton.config import NewtonConfig + +logger = logging.getLogger(__name__) + + +class NewtonSimulation(SimEngine): + """GPU-native simulation backend built on NVIDIA Warp + Newton. + + Implements the ``SimEngine`` ABC, providing the same programmatic + interface as MuJoCo/Isaac backends while enabling: + + - 4096+ parallel environments on a single GPU + - 7 solver backends (MuJoCo, Featherstone, XPBD, VBD, …) + - Differentiable simulation via ``wp.Tape`` + - CUDA graph capture for minimal Python overhead + - Soft-body / cloth / MPM simulation + + Parameters + ---------- + config : NewtonConfig | None + Backend configuration. Uses sensible defaults if ``None``. + + Examples + -------- + >>> from strands_robots.simulation.newton import NewtonSimulation, NewtonConfig + >>> config = NewtonConfig(num_envs=1, solver="mujoco", device="cpu") + >>> sim = NewtonSimulation(config=config) + >>> sim.create_world() # doctest: +SKIP + """ + + def __init__(self, config: NewtonConfig | None = None, **kwargs: Any) -> None: + if config is None: + # Accept factory kwargs (num_envs=…, solver=…) as NewtonConfig fields + config_kwargs = {k: v for k, v in kwargs.items() if k in NewtonConfig.__dataclass_fields__} + config = NewtonConfig(**config_kwargs) + + self._config = config + + # Warp / Newton modules — populated by _lazy_init() + self._wp: Any = None + self._newton: Any = None + + # Core simulation objects — populated by create_world() / _finalize_model() + self._builder: Any = None + self._model: Any = None + self._solver: Any = None + self._state_0: Any = None + self._state_1: Any = None + self._control: Any = None + self._contacts: Any = None + self._collision_pipeline: Any = None + self._renderer: Any = None + + # Entity tracking + self._robots: dict[str, dict[str, Any]] = {} + self._objects: dict[str, dict[str, Any]] = {} + self._sensors: dict[str, Any] = {} + + # State flags + self._world_created: bool = False + self._replicated: bool = False + self._step_count: int = 0 + self._sim_time: float = 0.0 + + # Pending action buffer for send_action() → step() pattern + self._pending_actions: dict[str, Any] | None = None + + logger.info( + "NewtonSimulation created — solver=%s, device=%s, num_envs=%d", + config.solver, + config.device, + config.num_envs, + ) + + # ------------------------------------------------------------------ + # Private helpers + # ------------------------------------------------------------------ + + def _lazy_init(self) -> None: + """Import warp and newton on first use. + + Raises + ------ + ImportError + If ``warp-lang`` or ``newton-sim`` are not installed. + """ + if self._wp is not None: + return + + try: + import warp as wp + + self._wp = wp + except ImportError as exc: + raise ImportError( + "warp-lang is required for the Newton backend. Install with: pip install 'strands-robots[newton]'" + ) from exc + + try: + import newton + + self._newton = newton + except ImportError as exc: + raise ImportError( + "newton-sim is required for the Newton backend. Install with: pip install 'strands-robots[newton]'" + ) from exc + + # Initialize Warp runtime + try: + self._wp.init() + logger.info("Warp initialised on %r.", self._config.device) + except Exception as exc: + logger.warning( + "Warp init on %r failed (%s), falling back to 'cpu'.", + self._config.device, + exc, + ) + self._config.device = "cpu" + + def _ensure_world(self) -> None: + """Raise if ``create_world()`` has not been called.""" + if not self._world_created: + raise RuntimeError("World not created. Call create_world() first.") + + # ------------------------------------------------------------------ + # SimEngine — World lifecycle (required) + # ------------------------------------------------------------------ + + def create_world( + self, + timestep: float | None = None, + gravity: list[float] | None = None, + ground_plane: bool = True, + ) -> dict[str, Any]: + """Create a new Newton simulation world. + + Initialises the Warp runtime (lazy) and creates a + ``newton.ModelBuilder`` with the requested physics parameters. + + Parameters + ---------- + timestep : float | None + Override ``physics_dt`` from config. + gravity : list[float] | None + 3-element gravity vector. Defaults to ``[0, -9.81, 0]``. + ground_plane : bool + Whether to add a ground plane. + + Returns + ------- + dict + ``{"success": True, "world_info": {…}}`` on success. + """ + self._lazy_init() + + if timestep is not None: + self._config.physics_dt = timestep + + # Will be implemented in PR 2 + raise NotImplementedError("create_world() — coming in PR 2 (world lifecycle)") + + def destroy(self) -> dict[str, Any]: + """Destroy the simulation world and release GPU resources.""" + raise NotImplementedError("destroy() — coming in PR 2 (world lifecycle)") + + def reset(self) -> dict[str, Any]: + """Reset all environments to their initial state.""" + raise NotImplementedError("reset() — coming in PR 2 (world lifecycle)") + + def step(self, n_steps: int = 1) -> dict[str, Any]: + """Advance physics by *n_steps* frames. + + Each frame applies ``self._config.substeps`` sub-steps of the + configured solver. Pending actions from ``send_action()`` are + applied at the start of each frame. + """ + raise NotImplementedError("step() — coming in PR 3 (step/action/observation)") + + def get_state(self) -> dict[str, Any]: + """Return a summary of the current simulation state.""" + raise NotImplementedError("get_state() — coming in PR 3 (step/action/observation)") + + # ------------------------------------------------------------------ + # SimEngine — Robot management (required) + # ------------------------------------------------------------------ + + def add_robot( + self, + name: str, + urdf_path: str | None = None, + data_config: str | None = None, + position: list[float] | None = None, + orientation: list[float] | None = None, + ) -> dict[str, Any]: + """Add a robot from URDF, MJCF, USD, or procedural definition. + + Supports automatic asset resolution and procedural fallback + for known robots (so100, koch, g1, go2). + """ + raise NotImplementedError("add_robot() — coming in PR 2 (world lifecycle)") + + def remove_robot(self, name: str) -> dict[str, Any]: + """Remove a robot from the simulation.""" + raise NotImplementedError("remove_robot() — coming in PR 4 (object management)") + + # ------------------------------------------------------------------ + # SimEngine — Object management (required) + # ------------------------------------------------------------------ + + def add_object( + self, + name: str, + shape: str = "box", + position: list[float] | None = None, + orientation: list[float] | None = None, + size: list[float] | None = None, + color: list[float] | None = None, + mass: float = 0.1, + is_static: bool = False, + mesh_path: str | None = None, + **kwargs: Any, + ) -> dict[str, Any]: + """Add a primitive or mesh object to the scene.""" + raise NotImplementedError("add_object() — coming in PR 4 (object management)") + + def remove_object(self, name: str) -> dict[str, Any]: + """Remove an object from the scene.""" + raise NotImplementedError("remove_object() — coming in PR 4 (object management)") + + # ------------------------------------------------------------------ + # SimEngine — Observation / Action (required) + # ------------------------------------------------------------------ + + def get_observation( + self, + robot_name: str | None = None, + camera_name: str | None = None, + ) -> dict[str, Any]: + """Get observation from simulation. + + Returns joint positions, velocities, body transforms, and + optionally a camera image. + """ + raise NotImplementedError("get_observation() — coming in PR 3 (step/action/observation)") + + def send_action( + self, + action: dict[str, Any], + robot_name: str | None = None, + n_substeps: int = 1, + ) -> None: + """Buffer an action for the next ``step()`` call. + + Actions are stored in ``self._pending_actions`` and applied + at the start of each physics frame in ``step()``. + """ + raise NotImplementedError("send_action() — coming in PR 3 (step/action/observation)") + + # ------------------------------------------------------------------ + # SimEngine — Rendering (required) + # ------------------------------------------------------------------ + + def render( + self, + camera_name: str = "default", + width: int | None = None, + height: int | None = None, + ) -> dict[str, Any]: + """Render an RGB frame from the specified camera.""" + raise NotImplementedError("render() — coming in PR 4 (rendering)") + + # ------------------------------------------------------------------ + # SimEngine — Optional overrides + # ------------------------------------------------------------------ + + def load_scene(self, scene_path: str) -> dict[str, Any]: + """Load a scene from URDF/MJCF/USD file.""" + raise NotImplementedError("load_scene() — coming in PR 2 (world lifecycle)") + + def run_policy( + self, + robot_name: str, + policy_provider: str = "mock", + **kwargs: Any, + ) -> dict[str, Any]: + """Run a policy loop in the simulation.""" + raise NotImplementedError("run_policy() — coming in PR 6 (advanced features)") + + def get_contacts(self) -> dict[str, Any]: + """Get contact information from the collision pipeline.""" + raise NotImplementedError("get_contacts() — coming in PR 4 (object management)") + + def cleanup(self) -> None: + """Release all GPU resources.""" + if self._renderer is not None: + try: + self._renderer = None + except Exception: + pass + self._model = None + self._solver = None + self._state_0 = None + self._state_1 = None + self._builder = None + self._world_created = False + logger.debug("NewtonSimulation cleanup complete.") + + # ------------------------------------------------------------------ + # Newton-specific extensions (NOT in SimEngine ABC) + # ------------------------------------------------------------------ + + def replicate(self, num_envs: int | None = None) -> dict[str, Any]: + """Clone the world into multiple parallel environments. + + Parameters + ---------- + num_envs : int | None + Number of environments. Defaults to ``config.num_envs``. + + Returns + ------- + dict + Replication result with ``num_envs`` and timing info. + """ + raise NotImplementedError("replicate() — coming in PR 5 (multi-env)") + + def run_diffsim( + self, + num_steps: int, + loss_fn: Any, + optimize_params: str, + lr: float = 0.02, + iterations: int = 200, + ) -> dict[str, Any]: + """Run a differentiable simulation optimisation loop. + + Uses ``wp.Tape`` for automatic differentiation through the + physics simulation. + """ + raise NotImplementedError("run_diffsim() — coming in PR 6 (advanced features)") + + def solve_ik( + self, + robot_name: str, + target_position: list[float], + target_orientation: list[float] | None = None, + ) -> dict[str, Any]: + """Solve inverse kinematics for a robot end-effector.""" + raise NotImplementedError("solve_ik() — coming in PR 6 (advanced features)") + + def add_cloth(self, name: str, **kwargs: Any) -> dict[str, Any]: + """Add a cloth body to the simulation.""" + raise NotImplementedError("add_cloth() — coming in PR 5 (soft bodies)") + + def add_cable(self, name: str, **kwargs: Any) -> dict[str, Any]: + """Add a cable body to the simulation.""" + raise NotImplementedError("add_cable() — coming in PR 5 (soft bodies)") + + def add_particles(self, name: str, **kwargs: Any) -> dict[str, Any]: + """Add MPM particles (granular/fluid) to the simulation.""" + raise NotImplementedError("add_particles() — coming in PR 5 (soft bodies)") + + def add_sensor(self, name: str, kind: str, **kwargs: Any) -> dict[str, Any]: + """Add a sensor (contact, IMU, or camera).""" + raise NotImplementedError("add_sensor() — coming in PR 6 (advanced features)") + + def read_sensor(self, name: str) -> dict[str, Any]: + """Read the latest value from a sensor.""" + raise NotImplementedError("read_sensor() — coming in PR 6 (advanced features)") + + def enable_dual_solver( + self, + articulated: str = "mujoco", + soft: str = "vbd", + ) -> None: + """Enable dual-solver mode (rigid + cloth solvers).""" + raise NotImplementedError("enable_dual_solver() — coming in PR 5 (soft bodies)") + + def reset_envs(self, env_ids: list[int]) -> dict[str, Any]: + """Reset specific environment IDs (Newton extension beyond ABC). + + The ABC ``reset()`` resets all envs. This method allows + selective per-env resets for RL training. + """ + raise NotImplementedError("reset_envs() — coming in PR 2 (world lifecycle)") diff --git a/tests/simulation/__init__.py b/tests/simulation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/simulation/newton/__init__.py b/tests/simulation/newton/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/simulation/newton/test_config.py b/tests/simulation/newton/test_config.py new file mode 100644 index 0000000..8b76593 --- /dev/null +++ b/tests/simulation/newton/test_config.py @@ -0,0 +1,145 @@ +"""Tests for NewtonConfig validation. + +Tests config construction, validation, defaults, and edge cases. +No GPU or heavy dependencies required — config is a pure dataclass. +""" + +from __future__ import annotations + +import pytest + +from strands_robots.simulation.newton.config import ( + BROAD_PHASE_OPTIONS, + RENDER_BACKENDS, + SOLVER_MAP, + NewtonConfig, +) + + +class TestNewtonConfigDefaults: + """Verify default config values match documented expectations.""" + + def test_default_values(self) -> None: + config = NewtonConfig() + assert config.num_envs == 1 + assert config.device == "cuda:0" + assert config.solver == "mujoco" + assert config.physics_dt == 0.005 + assert config.substeps == 1 + assert config.render_backend == "none" + assert config.enable_cuda_graph is False + assert config.enable_differentiable is False + assert config.broad_phase == "sap" + + def test_soft_contact_defaults(self) -> None: + config = NewtonConfig() + assert config.soft_contact_margin == 0.5 + assert config.soft_contact_ke == 10000.0 + assert config.soft_contact_kd == 10.0 + assert config.soft_contact_mu == 0.5 + assert config.soft_contact_restitution == 0.0 + + +class TestNewtonConfigValidSolvers: + """All 7 solvers should be accepted.""" + + @pytest.mark.parametrize("solver", list(SOLVER_MAP.keys())) + def test_valid_solver(self, solver: str) -> None: + config = NewtonConfig(solver=solver) + assert config.solver == solver + + +class TestNewtonConfigValidRenderBackends: + """All render backends should be accepted.""" + + @pytest.mark.parametrize("backend", sorted(RENDER_BACKENDS)) + def test_valid_render_backend(self, backend: str) -> None: + config = NewtonConfig(render_backend=backend) + assert config.render_backend == backend + + +class TestNewtonConfigValidBroadPhase: + """All broad-phase options should be accepted.""" + + @pytest.mark.parametrize("bp", sorted(BROAD_PHASE_OPTIONS)) + def test_valid_broad_phase(self, bp: str) -> None: + config = NewtonConfig(broad_phase=bp) + assert config.broad_phase == bp + + +class TestNewtonConfigInvalid: + """Invalid config values must raise immediately — fail-fast.""" + + def test_invalid_solver(self) -> None: + with pytest.raises(ValueError, match="Unknown solver"): + NewtonConfig(solver="nonexistent") + + def test_invalid_render_backend(self) -> None: + with pytest.raises(ValueError, match="Unknown render_backend"): + NewtonConfig(render_backend="metal") + + def test_invalid_broad_phase(self) -> None: + with pytest.raises(ValueError, match="Unknown broad_phase"): + NewtonConfig(broad_phase="octree") + + def test_zero_physics_dt(self) -> None: + with pytest.raises(ValueError, match="physics_dt must be positive"): + NewtonConfig(physics_dt=0.0) + + def test_negative_physics_dt(self) -> None: + with pytest.raises(ValueError, match="physics_dt must be positive"): + NewtonConfig(physics_dt=-0.001) + + def test_zero_num_envs(self) -> None: + with pytest.raises(ValueError, match="num_envs must be >= 1"): + NewtonConfig(num_envs=0) + + def test_negative_num_envs(self) -> None: + with pytest.raises(ValueError, match="num_envs must be >= 1"): + NewtonConfig(num_envs=-1) + + +class TestNewtonConfigCustom: + """Non-default config combinations.""" + + def test_gpu_training_config(self) -> None: + """4096-env GPU training scenario.""" + config = NewtonConfig( + num_envs=4096, + device="cuda:0", + solver="mujoco", + physics_dt=1.0 / 60.0, + substeps=4, + enable_cuda_graph=True, + ) + assert config.num_envs == 4096 + assert config.enable_cuda_graph is True + assert config.substeps == 4 + + def test_diffsim_config(self) -> None: + """Differentiable simulation scenario.""" + config = NewtonConfig( + solver="semi_implicit", + enable_differentiable=True, + enable_cuda_graph=False, # CUDA graphs + grad tracking conflict + ) + assert config.enable_differentiable is True + assert config.enable_cuda_graph is False + + def test_cpu_fallback(self) -> None: + config = NewtonConfig(device="cpu") + assert config.device == "cpu" + + +class TestSolverMap: + """SOLVER_MAP contains expected entries.""" + + def test_seven_solvers(self) -> None: + assert len(SOLVER_MAP) == 7 + + def test_mujoco_solver_class_name(self) -> None: + assert SOLVER_MAP["mujoco"] == "SolverMuJoCo" + + def test_all_solver_class_names_start_with_solver(self) -> None: + for name, cls_name in SOLVER_MAP.items(): + assert cls_name.startswith("Solver"), f"{name} → {cls_name}" diff --git a/tests/simulation/newton/test_factory.py b/tests/simulation/newton/test_factory.py new file mode 100644 index 0000000..ece67bf --- /dev/null +++ b/tests/simulation/newton/test_factory.py @@ -0,0 +1,90 @@ +"""Tests for Newton backend factory registration and alias resolution. + +Verifies that ``create_simulation("newton")`` and aliases ("warp", "wp") +correctly resolve to ``NewtonSimulation``. No GPU required — tests +only exercise the factory + import machinery, not the simulation itself. +""" + +from __future__ import annotations + +import pytest + +from strands_robots.simulation.base import SimEngine +from strands_robots.simulation.factory import ( + _BUILTIN_ALIASES, + _BUILTIN_BACKENDS, + _resolve_name, + create_simulation, + list_backends, +) +from strands_robots.simulation.newton.simulation import NewtonSimulation + + +class TestNewtonRegistration: + """Newton must be in the built-in backend registry.""" + + def test_newton_in_builtin_backends(self) -> None: + assert "newton" in _BUILTIN_BACKENDS + + def test_newton_module_path(self) -> None: + mod, cls = _BUILTIN_BACKENDS["newton"] + assert mod == "strands_robots.simulation.newton.simulation" + assert cls == "NewtonSimulation" + + def test_warp_alias(self) -> None: + assert _BUILTIN_ALIASES["warp"] == "newton" + + def test_wp_alias(self) -> None: + assert _BUILTIN_ALIASES["wp"] == "newton" + + +class TestNewtonInListBackends: + """list_backends() must include Newton and its aliases.""" + + def test_newton_in_list(self) -> None: + backends = list_backends() + assert "newton" in backends + + def test_warp_in_list(self) -> None: + backends = list_backends() + assert "warp" in backends + + def test_wp_in_list(self) -> None: + backends = list_backends() + assert "wp" in backends + + +class TestNewtonAliasResolution: + """Alias resolution must map to canonical 'newton'.""" + + @pytest.mark.parametrize("alias", ["newton", "warp", "wp"]) + def test_resolve_to_newton(self, alias: str) -> None: + assert _resolve_name(alias) == "newton" + + +class TestCreateNewtonSimulation: + """create_simulation("newton") must return a NewtonSimulation instance.""" + + def test_create_newton(self) -> None: + sim = create_simulation("newton") + assert isinstance(sim, NewtonSimulation) + assert isinstance(sim, SimEngine) + + def test_create_warp_alias(self) -> None: + sim = create_simulation("warp") + assert isinstance(sim, NewtonSimulation) + + def test_create_wp_alias(self) -> None: + sim = create_simulation("wp") + assert isinstance(sim, NewtonSimulation) + + def test_create_with_kwargs(self) -> None: + """Factory kwargs should flow through to NewtonConfig.""" + sim = create_simulation("newton", num_envs=64, solver="xpbd") + assert isinstance(sim, NewtonSimulation) + assert sim._config.num_envs == 64 + assert sim._config.solver == "xpbd" + + def test_create_unknown_backend_raises(self) -> None: + with pytest.raises(ValueError, match="Unknown simulation backend"): + create_simulation("nonexistent_engine") diff --git a/tests/simulation/newton/test_lazy_import.py b/tests/simulation/newton/test_lazy_import.py new file mode 100644 index 0000000..691efea --- /dev/null +++ b/tests/simulation/newton/test_lazy_import.py @@ -0,0 +1,60 @@ +"""Tests for Newton backend lazy-import behaviour. + +Verifies that importing ``strands_robots.simulation.newton`` does NOT +trigger import of ``warp`` or ``newton`` (heavy GPU dependencies). +This is critical for keeping ``import strands_robots`` fast and +usable on machines without CUDA. +""" + +from __future__ import annotations + +import importlib +import sys + + +class TestLazyImport: + """Importing the newton sub-package must not load warp or newton.""" + + def test_import_newton_package_does_not_load_warp(self) -> None: + """``import strands_robots.simulation.newton`` must be fast/light.""" + # Clear any cached imports so we get a clean test + mods_to_clear = [k for k in sys.modules if k.startswith("strands_robots.simulation.newton")] + for m in mods_to_clear: + del sys.modules[m] + + # Import the package + importlib.import_module("strands_robots.simulation.newton") + + # warp and newton must NOT have been imported + assert "warp" not in sys.modules, "warp was eagerly imported" + assert "newton" not in sys.modules, "newton was eagerly imported" + + def test_import_config_does_not_load_warp(self) -> None: + """Config is a pure dataclass — no GPU deps.""" + mods_to_clear = [k for k in sys.modules if k.startswith("strands_robots.simulation.newton")] + for m in mods_to_clear: + del sys.modules[m] + + from strands_robots.simulation.newton.config import NewtonConfig # noqa: F401 + + assert "warp" not in sys.modules + assert "newton" not in sys.modules + + def test_newton_config_available_directly(self) -> None: + """NewtonConfig must be importable from the package __init__.""" + from strands_robots.simulation.newton import NewtonConfig + + config = NewtonConfig() + assert config.solver == "mujoco" + + def test_factory_import_does_not_load_warp(self) -> None: + """create_simulation() import must not trigger Newton deps.""" + mods_to_clear = [k for k in sys.modules if k.startswith("strands_robots.simulation.newton")] + for m in mods_to_clear: + del sys.modules[m] + + from strands_robots.simulation import list_backends # noqa: F401 + + # Factory knows about newton but hasn't imported it yet + assert "warp" not in sys.modules + assert "newton" not in sys.modules diff --git a/tests/simulation/newton/test_simulation.py b/tests/simulation/newton/test_simulation.py new file mode 100644 index 0000000..c453911 --- /dev/null +++ b/tests/simulation/newton/test_simulation.py @@ -0,0 +1,195 @@ +"""Tests for NewtonSimulation class behaviour. + +Tests the skeleton class — instantiation, cleanup, SimEngine conformance, +and that unimplemented methods raise NotImplementedError with clear messages. +No GPU required. +""" + +from __future__ import annotations + +import pytest + +from strands_robots.simulation.base import SimEngine +from strands_robots.simulation.newton.config import NewtonConfig +from strands_robots.simulation.newton.simulation import NewtonSimulation + + +class TestNewtonSimulationInstantiation: + """NewtonSimulation can be instantiated without GPU.""" + + def test_default_config(self) -> None: + sim = NewtonSimulation() + assert isinstance(sim, SimEngine) + assert sim._config.solver == "mujoco" + assert sim._config.num_envs == 1 + + def test_custom_config(self) -> None: + config = NewtonConfig(num_envs=4096, solver="xpbd", device="cpu") + sim = NewtonSimulation(config=config) + assert sim._config.num_envs == 4096 + assert sim._config.solver == "xpbd" + + def test_kwargs_flow_to_config(self) -> None: + """Factory kwargs should populate config.""" + sim = NewtonSimulation(num_envs=128, solver="semi_implicit") + assert sim._config.num_envs == 128 + assert sim._config.solver == "semi_implicit" + + def test_initial_state_flags(self) -> None: + sim = NewtonSimulation() + assert sim._world_created is False + assert sim._replicated is False + assert sim._step_count == 0 + assert sim._sim_time == 0.0 + assert sim._pending_actions is None + assert sim._robots == {} + assert sim._objects == {} + + def test_is_simengine_subclass(self) -> None: + assert issubclass(NewtonSimulation, SimEngine) + + +class TestNewtonSimulationContextManager: + """Context manager protocol (__enter__/__exit__) works.""" + + def test_context_manager(self) -> None: + with NewtonSimulation() as sim: + assert isinstance(sim, NewtonSimulation) + # cleanup should have been called (no error) + + +class TestNewtonSimulationCleanup: + """cleanup() should be safe to call multiple times.""" + + def test_cleanup_idempotent(self) -> None: + sim = NewtonSimulation() + sim.cleanup() + sim.cleanup() # second call should not raise + + +class TestNewtonSimulationStubs: + """All ABC methods must raise NotImplementedError with clear messages. + + These are skeleton stubs — real implementations come in subsequent PRs. + Tests verify the fail-fast contract: callers get clear errors about + what's not yet available. + """ + + @pytest.fixture() + def sim(self) -> NewtonSimulation: + return NewtonSimulation() + + # --- World lifecycle --- + + def test_create_world_not_implemented(self, sim: NewtonSimulation) -> None: + """create_world requires warp, which is not installed in unit tests.""" + with pytest.raises((NotImplementedError, ImportError)): + sim.create_world() + + def test_destroy_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="destroy"): + sim.destroy() + + def test_reset_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="reset"): + sim.reset() + + def test_step_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="step"): + sim.step() + + def test_get_state_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="get_state"): + sim.get_state() + + # --- Robot management --- + + def test_add_robot_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="add_robot"): + sim.add_robot("so100") + + def test_remove_robot_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="remove_robot"): + sim.remove_robot("so100") + + # --- Object management --- + + def test_add_object_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="add_object"): + sim.add_object("cube") + + def test_remove_object_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="remove_object"): + sim.remove_object("cube") + + # --- Observation / Action --- + + def test_get_observation_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="get_observation"): + sim.get_observation("so100") + + def test_send_action_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="send_action"): + sim.send_action({"joint_0": 0.5}) + + # --- Rendering --- + + def test_render_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="render"): + sim.render() + + # --- Optional overrides --- + + def test_load_scene_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="load_scene"): + sim.load_scene("scene.usd") + + def test_run_policy_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="run_policy"): + sim.run_policy("so100") + + def test_get_contacts_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="get_contacts"): + sim.get_contacts() + + # --- Newton extensions --- + + def test_replicate_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="replicate"): + sim.replicate(4096) + + def test_run_diffsim_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="run_diffsim"): + sim.run_diffsim(100, lambda s: 0.0, "velocity") + + def test_solve_ik_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="solve_ik"): + sim.solve_ik("so100", [0.3, 0, 0.2]) + + def test_add_cloth_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="add_cloth"): + sim.add_cloth("cloth_0") + + def test_add_cable_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="add_cable"): + sim.add_cable("cable_0") + + def test_add_particles_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="add_particles"): + sim.add_particles("fluid_0") + + def test_add_sensor_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="add_sensor"): + sim.add_sensor("imu_0", "imu") + + def test_read_sensor_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="read_sensor"): + sim.read_sensor("imu_0") + + def test_enable_dual_solver_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="enable_dual_solver"): + sim.enable_dual_solver() + + def test_reset_envs_not_implemented(self, sim: NewtonSimulation) -> None: + with pytest.raises(NotImplementedError, match="reset_envs"): + sim.reset_envs([0, 1, 2])