cuju 0.1.0__py3-none-win_amd64.whl

cuju/__init__.py

@@ -0,0 +1,142 @@
+# Copyright 2026 The cuju Authors
+# Copyright 2019 Google LLC
+# SPDX-License-Identifier: Apache-2.0
+
+"""Batched Google Research Football (GRF) environments on the GPU or the CPU, bit for bit.
+
+`BatchEnvironment` steps many environments per call, and its outputs are PyTorch tensors that share
+its buffers. The guide and the reference are in the repository's `docs/` (`index.md`, `api.md`).
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+from typing import ClassVar
+
+from ._batch.assets import bundled_assets, bundled_scenarios
+from ._batch.environment import BatchEnvironment, memory_estimate
+from ._batch.library import library_path
+from ._batch.memory import _MemoryEstimate
+from ._version import __version__
+from .errors import (
+    ActionRangeError,
+    ArgumentError,
+    ChoiceError,
+    ClosedError,
+    ConfigurationError,
+    DeviceError,
+    EnvironmentIndexError,
+    Error,
+    LibraryError,
+    MissingFileError,
+    StateError,
+    UndefinedPathError,
+)
+
+__all__ = [
+    'ActionRangeError',
+    'ArgumentError',
+    'BatchEnvironment',
+    'ChoiceError',
+    'ClosedError',
+    'ConfigurationError',
+    'DeviceError',
+    'EnvironmentIndexError',
+    'Error',
+    'LibraryError',
+    'MissingFileError',
+    'MultiDeviceBatch',
+    'StateError',
+    'UndefinedPathError',
+    '__version__',
+    'bundled_assets',
+    'bundled_scenarios',
+    'library_path',
+    'memory_estimate',
+]
+
+# The names keep `cuju` as their module, so that their reprs and pickles name the public path.
+for _name in (
+    BatchEnvironment,
+    bundled_assets,
+    bundled_scenarios,
+    library_path,
+    memory_estimate,
+    _MemoryEstimate,
+):
+    _name.__module__ = __name__
+del _name
+
+# The digests and action logs build on `BatchEnvironment`.
+from .action_log import ActionLog, Digest, digest, replay  # noqa: E402
+
+__all__ += ['ActionLog', 'Digest', 'digest', 'replay']
+
+# Last, since it builds on `BatchEnvironment`.
+from .multi_device import MultiDeviceBatch  # noqa: E402
+
+
+class _RegisteringLoader:
+    """Runs the loader of gymnasium or gym, then registers GRF's ids with the module."""
+
+    def __init__(self, loader, name):
+        self._loader = loader
+        self._name = name
+
+    def create_module(self, spec):
+        return self._loader.create_module(spec)
+
+    def exec_module(self, module):
+        self._loader.exec_module(module)
+        _RegistrationFinder.waiting.discard(self._name)
+        _register_ids(self._name)
+
+    def __getattr__(self, name):
+        return getattr(self._loader, name)
+
+
+class _RegistrationFinder:
+    """Hands gymnasium and gym, when imported, a loader that then registers GRF's ids."""
+
+    waiting: ClassVar[set[str]] = set()
+    # The names being looked up by the other finders, which this one then leaves to them.
+    finding: ClassVar[set[str]] = set()
+
+    def find_spec(self, name, path=None, target=None):
+        if name not in self.waiting or name in self.finding:
+            return None
+        self.finding.add(name)
+        try:
+            spec = importlib.util.find_spec(name)
+        finally:
+            self.finding.discard(name)
+        if spec is not None and spec.loader is not None:
+            spec.loader = _RegisteringLoader(spec.loader, name)
+        return spec
+
+
+def _register_ids(module):
+    # A `cuju.env` that is still importing registers the ids itself once it is done.
+    try:
+        from .env import _register_with
+    except ImportError:
+        return
+    _register_with(module)
+
+
+def _register_gymnasium_ids():
+    """Registers GRF's ids with gymnasium and gym before 0.26, as importing `gfootball` does.
+
+    Each gets them once it is imported, so that `import cuju` imports neither.
+    """
+    for module in ('gymnasium', 'gym'):
+        if module in sys.modules:
+            _register_ids(module)
+        else:
+            _RegistrationFinder.waiting.add(module)
+    if not any(isinstance(finder, _RegistrationFinder) for finder in sys.meta_path):
+        sys.meta_path.insert(0, _RegistrationFinder())
+
+
+_register_gymnasium_ids()

cuju/_batch/__init__.py

@@ -0,0 +1,4 @@
+# Copyright 2026 The cuju Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""The implementation of `BatchEnvironment`, which `cuju` exports."""

cuju/_batch/allocator.py

@@ -0,0 +1,52 @@
+# Copyright 2026 The cuju Authors
+# SPDX-License-Identifier: Apache-2.0
+"""The device memory of a batch through PyTorch's caching allocator."""
+
+from __future__ import annotations
+
+import ctypes
+
+import torch
+
+# The C signatures of `allocate` and `release` in `cuju_configuration`.
+ALLOCATE = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.c_size_t, ctypes.c_void_p)
+RELEASE = ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_void_p)
+
+
+class TorchAllocator:
+    """The `allocate` and `release` functions that the library calls, over torch's allocator.
+
+    The memory comes from the default stream of the device, which the library's copies and fills
+    use, and the library releases it only once the device is idle, so torch may hand it out again
+    at once.
+
+    Attributes:
+        allocate, release: The C functions, which must outlive the batch.
+        errors: The exceptions that torch raised in `allocate`, such as `torch.OutOfMemoryError`.
+    """
+
+    def __init__(self, device: torch.device) -> None:
+        stream = torch.cuda.default_stream(device)
+        errors: list[BaseException] = []
+
+        def allocate(size, _context):
+            try:
+                return torch.cuda.caching_allocator_alloc(size, device, stream)
+            # An exception cannot cross the library, so the batch raises it after the call.
+            except Exception as error:
+                errors.append(error)
+                return None
+
+        def release(pointer, _context):
+            torch.cuda.caching_allocator_delete(pointer)
+
+        self.allocate = ALLOCATE(allocate)
+        self.release = RELEASE(release)
+        self.errors = errors
+
+    def pointers(self) -> tuple[int, int]:
+        """Returns the addresses of `allocate` and `release` for `cuju_configuration`."""
+        return (
+            ctypes.cast(self.allocate, ctypes.c_void_p).value,
+            ctypes.cast(self.release, ctypes.c_void_p).value,
+        )

cuju/_batch/assets.py

@@ -0,0 +1,68 @@
+# Copyright 2026 The cuju Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""The asset bundles that the package ships, and the bundle that a name or path asks for."""
+
+from __future__ import annotations
+
+import os
+import re
+
+from ..errors import MissingFileError
+
+_DATA_DIRECTORY = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'data')
+
+
+def bundled_assets() -> list[str]:
+    """Returns the names of the asset bundles that the package ships, such as `5_vs_5_l4_r0`.
+
+    `BatchEnvironment` takes such a name `<scenario>_l<left agents>_r<right agents>` in place of a
+    path, with any number of agents of a shipped scenario.
+    """
+    directory = os.path.join(_DATA_DIRECTORY, 'assets')
+    if not os.path.isdir(directory):
+        return []
+    return sorted(name[: -len('.cjb')] for name in os.listdir(directory) if name.endswith('.cjb'))
+
+
+def _split_name(name):
+    """Returns the scenario, left and right agents of `<scenario>_l<left>_r<right>`, or `None`."""
+    match = re.fullmatch(r'(.+)_l(\d+)_r(\d+)', name)
+    return None if match is None else (match.group(1), int(match.group(2)), int(match.group(3)))
+
+
+def bundled_scenarios() -> list[str]:
+    """Returns the scenarios of the asset bundles that the package ships, such as `5_vs_5`."""
+    return sorted({split[0] for split in map(_split_name, bundled_assets()) if split})
+
+
+def _resolve_asset(asset: str) -> tuple[str, tuple[int, int]]:
+    """Returns the bundle path of a path or name, and the agents per side it asks for (-1 keeps)."""
+    if callable(asset):
+        raise TypeError(
+            'expected the path of a scenario file or of a bundle, found a function; pass the path '
+            'of its file, or build a bundle with cuju.scenarios.build_bundle'
+        )
+    # GRF: `create_environment` gives a scenario file or a bare name one left agent.
+    if os.path.isfile(asset):
+        return asset, (1, 0) if asset.endswith('.py') else (-1, -1)
+    if asset.endswith(('.py', '.cjb')) or os.sep in asset or '/' in asset:
+        raise MissingFileError(f'expected a scenario file or a bundle at {asset!r}, found no file')
+    directory = os.path.join(_DATA_DIRECTORY, 'assets')
+    shipped = os.path.join(directory, asset + '.cjb')
+    if os.path.isfile(shipped):
+        return shipped, (-1, -1)
+    split = _split_name(asset)
+    if split is None and asset in bundled_scenarios():
+        split = (asset, 1, 0)
+    if split:
+        scenario, left, right = split
+        # The bundle with the fewest agents serves any counts; the library sets up the agents anew.
+        bases = [name for name in bundled_assets() if (_split_name(name) or ('',))[0] == scenario]
+        if bases:
+            base = min(bases, key=lambda name: sum(_split_name(name)[1:]))
+            return os.path.join(directory, base + '.cjb'), (left, right)
+    raise MissingFileError(
+        f'expected the path of an asset bundle or `<scenario>_l<left>_r<right>` of one of '
+        f'{bundled_scenarios()}, found {asset!r}'
+    )