stable-envs 0.1.0__py3-none-any.whl

stable_envs/__about__.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+__title__ = "stable-envs"
+__version__ = "0.1.0"
+__summary__ = "A batched research-environment library with explicit state, dynamics, variation spaces, and fast state-driven CPU/GPU rendering."
+__author__ = "Randall Balestriero, Sami BuGhanem"
+__license__ = "MIT"
+
+__all__ = ["__title__", "__version__", "__summary__", "__author__", "__license__"]

stable_envs/__init__.py

@@ -0,0 +1,91 @@
+"""Batched environments with explicit state, variation and state-driven rendering.
+
+Environment classes are attributes named by their registry key, for example ``stable_envs.PushT``.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from . import registry as _registry
+from . import tuning
+from .__about__ import __version__  # noqa: F401
+from ._core.config import get_config, set  # noqa: A004,F401
+from .env import Env, State, StepResult, VariationLike  # noqa: F401
+from .registry import import_errors, info, list_envs, list_families  # noqa: F401
+from .registry import make as _make
+from .specs import CompositeSpec, Spec  # noqa: F401
+
+# Discover available families; report missing optional dependencies through import_errors().
+_registry.discover()
+
+
+def make(name: str, **kwargs) -> Env:
+    """Construct a registered environment by class name.
+
+    ``make("CartPole", num_envs=64)`` is the same as ``stable_envs.CartPole(num_envs=64)``.
+
+    Args:
+        name: Class name from ``list_envs``, not a ``reference_id`` such as ``"CartPole-v1"``.
+        **kwargs: Forwarded to the constructor, such as ``num_envs``, ``size``, ``device`` and ``seed``.
+
+    Raises:
+        KeyError: If ``name`` is not registered. ``import_errors`` shows families that didn't import.
+    """
+    return _make(name, **kwargs)
+
+
+def tune(env: Env, *, knobs=(), repeats: int = 5, write: bool = True, verbose: bool = True) -> dict[str, int]:
+    """Time the tuning knobs ``env`` reads and cache the fastest values.
+
+    Args:
+        env: Environment whose device, ``size`` and ``num_envs`` define the workload.
+        knobs: Knobs to sweep (all that ``env`` reads if empty).
+        repeats: Timed calls per value.
+        write: Save the results to the tuning cache.
+        verbose: Print the result table.
+
+    Returns:
+        A dict mapping each swept knob to its chosen value.
+
+    Raises:
+        ValueError: If ``knobs`` names an unknown knob.
+    """
+    from ._tune import tune as _tune
+
+    return _tune(env, knobs=knobs, repeats=repeats, write=write, echo=print if verbose else None)
+
+
+def __getattr__(name: str) -> Any:
+    """Resolve environment classes by canonical registry name."""
+    if name in _registry.REGISTRY:
+        return _registry.REGISTRY[name]
+    raise AttributeError(
+        f"module 'stable_envs' has no attribute {name!r}. Environments are looked up by class "
+        f"name: snv.list_envs() shows the {len(_registry.REGISTRY)} registered and "
+        f"snv.import_errors() shows the families that did not import."
+    )
+
+
+def __dir__() -> list[str]:
+    return sorted({*__all__, *_registry.REGISTRY})
+
+
+__all__ = [
+    "Env",
+    "State",
+    "StepResult",
+    "VariationLike",
+    "Spec",
+    "CompositeSpec",
+    "make",
+    "info",
+    "list_envs",
+    "list_families",
+    "import_errors",
+    "set",
+    "get_config",
+    "tuning",
+    "tune",
+    "__version__",
+]

stable_envs/__main__.py

@@ -0,0 +1,3 @@
+from .cli import main
+
+raise SystemExit(main())

stable_envs/_core/__init__.py

@@ -0,0 +1 @@
+"""Private shared helpers used across the stable_envs package."""

stable_envs/_core/cache.py

@@ -0,0 +1,12 @@
+"""The stable-envs cache directory, shared by the tuning cache and model downloads."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+
+def cache_dir() -> Path:
+    """Return the ``stable-envs`` cache root: ``$XDG_CACHE_HOME/stable-envs``, or ``~/.cache/stable-envs`` when unset."""
+    base = os.environ.get("XDG_CACHE_HOME") or str(Path.home() / ".cache")
+    return Path(base) / "stable-envs"

stable_envs/_core/config.py

@@ -0,0 +1,84 @@
+"""Global verbosity and default-device settings, and the package logger."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from loguru import logger as _loguru_logger
+
+from . import environ as _environ
+
+# Records carry extra["package"] == "stable_envs" for application-side filtering.
+logger = _loguru_logger.bind(package="stable_envs")
+
+_LEVELS = {
+    "TRACE": 5,
+    "DEBUG": 10,
+    "INFO": 20,
+    "SUCCESS": 25,
+    "WARNING": 30,
+    "ERROR": 40,
+    "CRITICAL": 50,
+    "SILENT": 100,
+}
+
+_DEFAULTS: dict[str, Any] = {"verbose": "WARNING", "default_device": "auto"}
+
+
+def _initial() -> dict[str, Any]:
+    return {**_DEFAULTS, "verbose": _environ.read("STABLE_ENVS_VERBOSE")}
+
+
+_CONFIG: dict[str, Any] = _initial()
+
+
+def set(**kwargs: Any) -> dict[str, Any]:  # noqa: A001 - the public spelling is ``snv.set``
+    """Update the log level or default device and return a copy of the settings.
+
+    Args:
+        **kwargs: ``verbose``, the lowest level logged (``"TRACE"``, ``"DEBUG"``, ``"INFO"``,
+            ``"SUCCESS"``, ``"WARNING"``, ``"ERROR"``, ``"CRITICAL"`` or ``"SILENT"``), and
+            ``default_device``, used when a constructor gets ``device=None``.
+
+    Raises:
+        KeyError: If any other key is passed.
+        ValueError: If ``verbose`` is not a known level.
+    """
+    unknown = [k for k in kwargs if k not in _DEFAULTS]
+    if unknown:
+        raise KeyError(
+            f"stable_envs.set: unknown config key(s) {unknown}; the keys are {sorted(_DEFAULTS)}. "
+            f"device, size, supersample and backend are environment constructor arguments."
+        )
+    if "verbose" in kwargs:
+        v = str(kwargs["verbose"]).upper()
+        if v not in _LEVELS:
+            raise ValueError(f"stable_envs.set: verbose must be one of {sorted(_LEVELS)}, got {kwargs['verbose']!r}")
+        kwargs["verbose"] = v
+    if "default_device" in kwargs and kwargs["default_device"] is not None:
+        kwargs["default_device"] = str(kwargs["default_device"])
+    _CONFIG.update(kwargs)
+    return dict(_CONFIG)
+
+
+def get_config() -> dict[str, Any]:
+    """Return a copy of the global settings, ``verbose`` and ``default_device``."""
+    return dict(_CONFIG)
+
+
+def reset_config() -> dict[str, Any]:
+    """Restore and return the default global settings, with verbosity from ``STABLE_ENVS_VERBOSE``."""
+    _CONFIG.clear()
+    _CONFIG.update(_initial())
+    return dict(_CONFIG)
+
+
+def log(level: str, message: str) -> None:
+    """Log ``message`` through the package logger when ``level`` meets the configured verbosity.
+
+    Args:
+        level: A level name such as ``"DEBUG"`` or ``"WARNING"``.
+        message: The text to log.
+    """
+    if _LEVELS[level.upper()] >= _LEVELS[_CONFIG["verbose"]]:
+        logger.opt(depth=1).log(level.upper(), message)

stable_envs/_core/dynamics.py

@@ -0,0 +1,17 @@
+"""Detect whether an environment implements dynamics, and the error raised when it does not."""
+
+
+def has_dynamics(env):
+    """Return whether ``env`` overrides the base ``_step``."""
+    base = getattr(env, "_se_dynamics_base", None)
+    if base is None:
+        return hasattr(env, "step")
+    return type(env)._step is not base._step
+
+
+def no_dynamics(env, what):
+    """Raise NotImplementedError for ``what`` on a render-only environment."""
+    raise NotImplementedError(
+        f"{type(env).__name__}.{what}: full simulator semantics are not implemented for this environment; "
+        "it renders supplied states only. Compact renderer rows cannot substitute for State."
+    )

stable_envs/_core/environ.py

@@ -0,0 +1,166 @@
+"""The registry of every ``STABLE_ENVS_*`` environment variable and the one function that reads them."""
+
+from __future__ import annotations
+
+import functools
+import os
+import re
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+PREFIX = "STABLE_ENVS_"
+
+
+@dataclass(frozen=True)
+class Variable:
+    """One environment variable: its name, value type, default, documentation and parser."""
+
+    name: str
+    kind: str
+    default: Any
+    doc: str
+    parse: Callable[[str], Any]
+
+
+VARIABLES: dict[str, Variable] = {}
+
+_SUFFIX = {"": 1, "K": 1 << 10, "M": 1 << 20, "G": 1 << 30, "T": 1 << 40}
+
+
+def parse_int(text: str) -> int | None:
+    """Parse an integer with an optional K, M, G or T suffix (powers of 1024); ``auto`` or empty is None."""
+    t = text.strip().replace("_", "")
+    if t.lower() in ("", "auto"):
+        return None
+    m = re.fullmatch(r"(\d+)\s*([KMGT]?)I?B?", t.upper())
+    if m is None:
+        raise ValueError("expected an integer, optionally with a K, M, G or T suffix, or 'auto'")
+    return int(m.group(1)) * _SUFFIX[m.group(2)]
+
+
+def parse_flag(text: str) -> bool:
+    """Parse ``1`` as True and ``0`` or empty as False."""
+    t = text.strip()
+    if t not in ("", "0", "1"):
+        raise ValueError("expected 1 or 0")
+    return t == "1"
+
+
+def parse_path(text: str) -> str | None:
+    """Return a non-empty path string, or None when empty."""
+    return text.strip() or None
+
+
+def register(name: str, kind: str, default: Any, doc: str, parse: Callable[[str], Any]) -> Variable:
+    """Add a variable to the registry; registering the same name twice keeps the first entry."""
+    if not name.startswith(PREFIX):
+        raise ValueError(f"{name} does not start with {PREFIX}")
+    return VARIABLES.setdefault(name, Variable(name, kind, default, doc, parse))
+
+
+@functools.lru_cache(maxsize=512)
+def _parsed(name: str, raw: str) -> Any:
+    try:
+        return VARIABLES[name].parse(raw)
+    except ValueError as e:
+        raise ValueError(f"{name}={raw!r}: {e}") from None
+
+
+def read(name: str) -> Any:
+    """Return the parsed value of variable ``name``, or its default when unset.
+
+    Raises:
+        KeyError: ``name`` is not registered.
+        ValueError: The value does not parse; the message names the variable.
+    """
+    var = VARIABLES[name]
+    raw = os.environ.get(name)
+    return var.default if raw is None else _parsed(name, raw)
+
+
+def variables() -> dict[str, Variable]:
+    """Return every registered variable, including those of modules not imported yet."""
+    from .. import tuning  # noqa: F401  (registers the knob variables)
+
+    return dict(sorted(VARIABLES.items()))
+
+
+def set_variables(environ=None) -> tuple[dict[str, str], list[str]]:
+    """Return the ``STABLE_ENVS_*`` variables set in ``environ`` and the names that are not registered."""
+    environ = os.environ if environ is None else environ
+    known = variables()
+    found = {k: v for k, v in sorted(environ.items()) if k.startswith(PREFIX)}
+    return found, [k for k in found if k not in known]
+
+
+def parse_verbosity(text: str) -> str:
+    """Parse a log level name such as ``INFO`` or ``SILENT``; an empty value is ``WARNING``."""
+    from .config import _LEVELS
+
+    t = text.strip().upper() or "WARNING"
+    if t not in _LEVELS:
+        raise ValueError(f"expected one of {sorted(_LEVELS)}")
+    return t
+
+
+def parse_dmc_render(text: str) -> str:
+    """Parse ``auto`` (or empty) or ``mjwarp``."""
+    t = text.strip() or "auto"
+    if t not in ("auto", "mjwarp"):
+        raise ValueError("expected auto or mjwarp")
+    return t
+
+
+def parse_json_object(text: str) -> dict | None:
+    """Parse a JSON object; an empty value is None."""
+    import json
+
+    t = text.strip()
+    if not t:
+        return None
+    try:
+        value = json.loads(t)
+    except json.JSONDecodeError as e:
+        raise ValueError(f"expected a JSON object: {e}") from None
+    if not isinstance(value, dict):
+        raise ValueError("expected a JSON object")
+    return value
+
+
+register(
+    "STABLE_ENVS_DMC_CUDA_RENDER",
+    "choice",
+    "auto",
+    "mjwarp makes backend='auto' fixed-camera DM Control environments on CUDA draw with MuJoCo Warp's ray tracer.",
+    parse_dmc_render,
+)
+register(
+    "STABLE_ENVS_DMC_WARP_OPTIONS",
+    "json",
+    None,
+    "A JSON object of default MuJoCo Warp render options for DM Control; render_options= overrides it.",
+    parse_json_object,
+)
+register(
+    "STABLE_ENVS_VERBOSE",
+    "level",
+    "WARNING",
+    "Default log level of the library, with the same values as snv.set(verbose=...). snv.set overrides it.",
+    parse_verbosity,
+)
+register(
+    "STABLE_ENVS_OGB_UNMEASURED",
+    "flag",
+    False,
+    "1 lets OGBench Cube and Scene cells without measured visual fidelity register, with placeholder "
+    "values, while measuring them. Never set it for normal use.",
+    parse_flag,
+)
+register(
+    "STABLE_ENVS_MSHUMAN_HOME",
+    "path",
+    None,
+    "Path to an existing MS-Human-700 checkout; unset downloads the pinned commit into the cache.",
+    parse_path,
+)

stable_envs/_core/mujoco_gl.py

@@ -0,0 +1,106 @@
+"""Configure GL before importing MuJoCo or dm_control.
+
+On macOS, bind MuJoCo to its native backend before importing dm_control with
+rendering disabled. This avoids GLFW initialization during discovery. Restore
+``glfw`` afterward for Gymnasium, which rejects ``cgl``. On other platforms,
+default to offscreen EGL and preserve explicit settings.
+"""
+
+from __future__ import annotations
+
+import importlib
+import importlib.util
+import os
+import sys
+
+_DARWIN_MUJOCO = frozenset({"cgl", "glfw"})
+_LINUX_MUJOCO = frozenset({"egl", "osmesa", "glfw"})
+GL_VARS = ("MUJOCO_GL", "PYOPENGL_PLATFORM", "MUJOCO_EGL_DEVICE_ID")
+#: The GL variables as the process received them, before configure_mujoco_gl fills defaults.
+INITIAL_GL_ENV = {name: os.environ.get(name) for name in GL_VARS}
+
+NO_RENDERER_FIX = (
+    "Headless Linux: install EGL (apt install libegl1; NVIDIA machines also need the driver's "
+    "libEGL_nvidia) and keep MUJOCO_GL=egl, or install OSMesa (apt install libosmesa6) and set "
+    "MUJOCO_GL=osmesa PYOPENGL_PLATFORM=osmesa. With a display, MUJOCO_GL=glfw. "
+    "`stable-envs doctor` tests each backend on this machine."
+)
+
+
+def selected_backend(environ=None, platform: str = sys.platform) -> str:
+    """Return the MUJOCO_GL backend MuJoCo binds when stable_envs imports it.
+
+    Unset on Linux, it follows a PYOPENGL_PLATFORM that names a MuJoCo backend, else EGL.
+    """
+    environ = os.environ if environ is None else environ
+    backend = environ.get("MUJOCO_GL")
+    if platform == "darwin":
+        return backend if backend in _DARWIN_MUJOCO else "cgl"
+    if backend is not None:
+        return backend
+    pyopengl = (environ.get("PYOPENGL_PLATFORM") or "").lower()
+    return pyopengl if pyopengl in _LINUX_MUJOCO else "egl"
+
+
+def configure_mujoco_gl() -> str:
+    """Configure import-time GL selection and return the resulting environment value."""
+    if sys.platform == "darwin":
+        if "mujoco" not in sys.modules:
+            os.environ["MUJOCO_GL"] = selected_backend()
+            try:
+                import mujoco  # noqa: F401
+            except Exception:  # noqa: BLE001 — optional extra; discovery records the miss
+                pass
+        have_dmc = importlib.util.find_spec("dm_control") is not None
+        if have_dmc and "dm_control._render" not in sys.modules:
+            os.environ["MUJOCO_GL"] = "disable"
+            try:
+                importlib.import_module("dm_control.mujoco")
+            except Exception:  # noqa: BLE001 — a backend that is not installed
+                return os.environ["MUJOCO_GL"]
+        # Gymnasium's OffScreenViewer only accepts glfw/egl/osmesa. dm_control is already
+        # imported (or absent), so restoring glfw does not re-enter GLFW init in this process.
+        os.environ["MUJOCO_GL"] = "glfw"
+        return os.environ["MUJOCO_GL"]
+    os.environ.setdefault("MUJOCO_GL", selected_backend())
+    return os.environ["MUJOCO_GL"]
+
+
+def _renderer_import_error() -> Exception | None:
+    # mujoco swallows the ImportError that leaves it without Renderer; importing again recovers it.
+    for name in ("mujoco.rendering.classic.renderer", "mujoco.renderer"):
+        try:
+            importlib.import_module(name)
+        except ModuleNotFoundError as e:
+            if e.name != name:
+                return e
+        except Exception as e:  # noqa: BLE001 — reported to the user
+            return e
+    return None
+
+
+def mujoco_renderer(model, *, height: int, width: int):
+    """Construct ``mujoco.Renderer``, failing with the GL setup to change when no OpenGL context can be made."""
+    import mujoco
+
+    backend = os.environ.get("MUJOCO_GL", "<unset>")
+    renderer = getattr(mujoco, "Renderer", None)
+    if renderer is None:
+        cause = _renderer_import_error()
+        detail = f"MuJoCo has no OpenGL renderer with MUJOCO_GL={backend}"
+    else:
+        try:
+            return renderer(model, height=height, width=width)
+        except ValueError:
+            raise
+        except Exception as e:  # noqa: BLE001 — re-raised with the fix
+            cause = e
+            detail = f"MuJoCo could not create an OpenGL context with MUJOCO_GL={backend}"
+    pyopengl = os.environ.get("PYOPENGL_PLATFORM")
+    if cause is not None and "PYOPENGL_PLATFORM" in str(cause) and pyopengl:
+        raise RuntimeError(
+            f"PYOPENGL_PLATFORM={pyopengl} conflicts with MUJOCO_GL={backend}: unset PYOPENGL_PLATFORM "
+            f"or set both to the same backend. {NO_RENDERER_FIX}"
+        ) from cause
+    reason = f" ({type(cause).__name__}: {cause})" if cause is not None else ""
+    raise RuntimeError(f"{detail}{reason}. {NO_RENDERER_FIX}") from cause