diff --git a/.agents/rules/architecture.md b/.agents/rules/architecture.md new file mode 100644 index 0000000..e1e71fa --- /dev/null +++ b/.agents/rules/architecture.md @@ -0,0 +1,20 @@ +# Architecture and Code Formatting Rules + +When writing or modifying code in this project, adhere strictly to the following rules: + +## 1. Core Frameworks & Tooling +- **PPO & CleanRL**: Proximal Policy Optimization (PPO) is the baseline algorithm. Use CleanRL as the starting framework for PPO, ensuring adaptation for continuous action spaces. +- **JAX / Flax**: All Artificial Neural Network (ANN) controller architectures must be implemented using Flax (neural networks in JAX). Ensure full compatibility with the JAX/Flax ecosystem. +- **MuJoCo**: The simulation environment uses a MuJoCo brittle star. Ensure that XML structures (sensors, actuators, joints, morphology) respect realistic constraints and adhere to the project requirements. + +## 2. Clean Code Principles +- **Naming Conventions**: Variables and functions must have consistent, intention-revealing names. A descriptive name is universally preferred over a short name with an explanatory comment. +- **Single Responsibility Function Design**: Functions must be modular. Minimize arguments and completely avoid boolean flag arguments that control execution behavior. +- **Commenting**: Code explains the "how". Comments are strictly reserved for explaining the "why". +- **No Commented-out Code**: The AI must **never** generate commented-out or dead code. Delete it using version control instead. +- **YAGNI & Complexity Management**: You Aren't Gonna Need It. Avoid premature optimization or unnecessary abstraction. Only introduce complexity with documented justification. Break large functions into small, testable blocks. +- **No Notebooks for Core Logic**: Jupyter Notebooks are explicitly forbidden for general software development as they discourage modularity. They should only be used for prototyping, tutorials, or post-processing analysis. + +## 3. Formatting & Linting +- **Ruff**: Output perfectly formatted code adhering to the Google style standard. Always format and lint the code using `ruff` (see `pyproject.toml` and `ruff.toml`). +- **Separation of Concerns**: Configuration code must be completely separated from implementation logic. Core logic must never be mixed with scripts or notebooks. diff --git a/.agents/rules/general.md b/.agents/rules/general.md new file mode 100644 index 0000000..c62f3ed --- /dev/null +++ b/.agents/rules/general.md @@ -0,0 +1,15 @@ +# General AI Agent Rules + +When assisting with this project, the AI Agent must strictly abide by these overarching operational rules: + +## 1. Scientific Integrity +- **No Hallucinations**: You must never fabricate results, hallucinate citations, or generate false empirical claims. Do not guess what happened if a process fails; rely strictly on outputs and logs. + +## 2. Agent Operational Constraints +- **Absolute Paths**: Always use absolute paths when making tool calls or reading/writing files. +- **Refactoring Guardrails**: Do not commence massive files/directory refactors or major system migrations without explicitly communicating the plan and asking for user clarification or approval first. +- **No Boilerplate Feedback**: The AI must not produce generic boilerplate summaries or overly generic advice. Ensure all outputs are completely contextual, robust, and well-reasoned. + +## 3. Communication +- **Artifacts and UI**: Use artifacts (like `implementation_plan.md` or `task.md`) appropriately for tracking design phases and updates. +- **Explicit Documenting**: If performing design choices, list them explicitly. Keep the output focused on the exact task. diff --git a/.agents/rules/git-workflow.md b/.agents/rules/git-workflow.md new file mode 100644 index 0000000..312e457 --- /dev/null +++ b/.agents/rules/git-workflow.md @@ -0,0 +1,25 @@ +# Git Workflow & Repository Structure Rules + +When performing Git operations and managing the repository layout, follow these rules: + +## 1. Committing Practices +- **Frequent & Small**: Produce small, logical commits instead of massive monolithic ones. +- **Conventional Commits**: Commit messages must adhere to the Conventional Commits specification (e.g., `feat: ...`, `fix: ...`, `refactor: ...`). +- **Single Functionality**: Each commit should relate to exactly one piece of functionality or distinct structural change. + +## 2. Branching & Merging +- **Branch `dev`**: The `dev` branch is the primary integration branch for pushing and merging code. +- **Branch `main`**: Only stable, finalized releases may be pushed to `main`. +- **Feature Branches**: Organize distinct work into logical feature branches when pushing to the remote server, maintaining an organized Git history. + +## 3. Artifact Management & Exclusions +- **LFS Only**: Data files, trained models, and large datasets must **never** be committed directly to Git. Ensure they are tracked with Git Large File Storage (LFS). **CRITICAL**: Every developer and AI agent must have `git-lfs` installed locally for the repository hooks to successfully pull these large files. Run `git lfs install` after cloning or setting up your environment. + +## 4. Repository Layout Strictness +Ensure generated code is meticulously placed in the correct directories: +- `src/` for algorithms, network designs, and core agent modules. +- `env/` for MuJoCo wrappers and environment definitions. +- `config/` for experiment configurations (using json, gin, or yaml). +- `experiments/` for executable scripts. +- `docs/` for ReadTheDocs or Doxygen documentation, and decision logs. +- `tests/` for unit tests and verification scripts. diff --git a/.agents/rules/method.md b/.agents/rules/method.md new file mode 100644 index 0000000..40e4b20 --- /dev/null +++ b/.agents/rules/method.md @@ -0,0 +1,14 @@ +# Methodology and Process Rules + +The agent must adhere to the following scientific and operational practices, focused on robustness and reproducibility: + +## 1. Scientific Context & Methodology +- **Research Focus**: Maintain focus on the project's objective: studying how controller modularity affects learning speed, coordination, and fault tolerance in brittle-star locomotion. +- **Hypothesis-Driven Design**: Base execution on clear hypotheses. Document all design decisions prior to implementation (in `/docs/decisions/` or via Artifacts/Plans). +- **Scaffolding Approach**: Start development with simple setups before scaling to complex environments and varying morphologies. +- **Value of Negative Results**: Understand that a controller failing to learn locomotion, when coupled with a thorough analysis of the failure, holds strong scientific value. Do not artificially force a positive result. + +## 2. Reproducibility Protection +- **Dependency Management (uv)**: This project strictly prefers `uv`. Do **not** manually modify the `uv.lock` file. Add dependencies via `uv add ` and sync environments via `uv sync --frozen` (or via devcontainers). +- **Consistent Initialization**: The AI must avoid hidden randomness. Ensure that all runs use consistent seed initialization. +- **Experiment Tracking**: Use Weights & Biases (wandb) for tracking and logging all experiments and parameters. Ensure every experiment-friendly parameter is properly logged. diff --git a/.agents/skills/lint/SKILL.md b/.agents/skills/lint/SKILL.md new file mode 100644 index 0000000..b5a70cf --- /dev/null +++ b/.agents/skills/lint/SKILL.md @@ -0,0 +1,24 @@ +--- +name: Lint and Format Code +description: Instructions for checking code style and formatting using Ruff. +--- + +# Skill: Lint and Format Code + +The goal of this skill is to enforce the project's adherence to the Google style standard and to maintain high code quality across the python source files. + +## Instructions + +1. **Format Code**: To automatically format all Python files according to the `ruff.toml` specifications: + ```bash + uv run ruff format . + ``` + +2. **Check for Lints / Auto-fix**: To check the repository for style violations and automatically fix safe corrections: + ```bash + uv run ruff check --fix . + ``` + +## Important Considerations +- The `ruff.toml` file at the root handles all lint and format configuration. Do not ignore configurations when applying fixes. +- If Ruff points out complex errors that cannot be auto-fixed, analyze the code and manually address the violations, prioritizing descriptive naming and adherence to the single-responsibility principle. diff --git a/.agents/skills/test/SKILL.md b/.agents/skills/test/SKILL.md new file mode 100644 index 0000000..93a939b --- /dev/null +++ b/.agents/skills/test/SKILL.md @@ -0,0 +1,25 @@ +--- +name: Run Tests +description: Instructions for executing the project test suite to verify code correctness. +--- + +# Skill: Run Tests + +The goal of this skill is to verify that the project is functioning correctly after development changes. + +## Instructions + +1. **Verify Environment**: The project operates dynamically with hardware acceleration (GPU via JAX) and uses `uv` for dependency/execution management. +2. **Execute Tests**: + - To run basic verification (e.g., checking JAX initialization and hardware detection), run the test script: + ```bash + uv run python -m tests.test_jax_init + ``` + - If a broader suite of modular tests is added (e.g., `pytest`), execute tests in the `tests/` directory with: + ```bash + uv run pytest tests/ + ``` + +## Important Considerations +- If you are running tests inside the local environment without a Devcontainer, verify whether `uv sync --frozen` (for CPU) or `uv sync --frozen --extra cuda` (for GPU) has been executed to avoid import errors. +- Do not run bare `python ...` without `uv run` locally, unless you are strictly operating inside a pre-activated `.venv` inside a Devcontainer. diff --git a/.commitlintrc.json b/.commitlintrc.json new file mode 100644 index 0000000..c30e5a9 --- /dev/null +++ b/.commitlintrc.json @@ -0,0 +1,3 @@ +{ + "extends": ["@commitlint/config-conventional"] +} diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile new file mode 100644 index 0000000..bebd64c --- /dev/null +++ b/.devcontainer/Dockerfile @@ -0,0 +1,30 @@ +FROM mcr.microsoft.com/devcontainers/python:3.12 + +# Install uv for dependency management +COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ + +# Install system dependencies for JAX, OpenGL, and Mujoco +RUN apt-get update && apt-get install -y --no-install-recommends \ + libgl1 \ + libosmesa6-dev \ + libglew-dev \ + libglfw3 \ + patchelf \ + && apt-get clean && rm -rf /var/lib/apt/lists/* + +# Set environment variables for hardware acceleration and rendering +ENV LD_LIBRARY_PATH=/usr/lib/nvidia +ENV NVIDIA_VISIBLE_DEVICES=all +ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility,graphics + +# Set the workspace path (defaults to /workspaces/project but can be overridden) +ARG WORKSPACE_PATH=/workspaces/project +ENV WORKSPACE_PATH=$WORKSPACE_PATH +ENV PATH="$WORKSPACE_PATH/.venv/bin:$PATH" + +# Ensure the workspace and .venv directories exist and are owned by the vscode user +# This prevents permission errors when the postCreateCommand runs uv sync +RUN mkdir -p $WORKSPACE_PATH/.venv && chown -R vscode:vscode $WORKSPACE_PATH + +USER vscode +WORKDIR $WORKSPACE_PATH diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json new file mode 100644 index 0000000..623526f --- /dev/null +++ b/.devcontainer/devcontainer.json @@ -0,0 +1,48 @@ +{ + "name": "Brittle Star JAX/CUDA", + "build": { + "dockerfile": "Dockerfile", + "context": "..", + "args": { + "WORKSPACE_PATH": "/workspaces/${localWorkspaceFolderBasename}" + } + }, + "workspaceFolder": "/workspaces/${localWorkspaceFolderBasename}", + "customizations": { + "vscode": { + "extensions": [ + "ms-python.python", + "charliermarsh.ruff", + "ms-python.vscode-pylance", + "ms-toolsai.jupyter" + ], + "settings": { + "python.defaultInterpreterPath": "${containerWorkspaceFolder}/.venv/bin/python", + "python.analysis.localRoot": "${containerWorkspaceFolder}", + "python.analysis.extraPaths": [ + "${containerWorkspaceFolder}/.venv/lib/python3.12/site-packages" + ] + } + }, + "jetbrains": { + "plugins": [ + "Pythonid", + "fleet.python", + "com.koxudaxi.ruff" + ] + } + }, + "remoteUser": "vscode", + "runArgs": [ + "--device", "nvidia.com/gpu=all" + ], + // Ensure the .venv persists using a named volume for performance and parity + "mounts": [ + "source=${localWorkspaceFolderBasename}-venv,target=${containerWorkspaceFolder}/.venv,type=volume" + ], + // Invoke the hardware-aware sync script + "postCreateCommand": "bash ./.devcontainer/post-create.sh", + "features": { + "ghcr.io/devcontainers/features/common-utils:1": {} + } +} diff --git a/.devcontainer/post-create.sh b/.devcontainer/post-create.sh new file mode 100644 index 0000000..576bad1 --- /dev/null +++ b/.devcontainer/post-create.sh @@ -0,0 +1,18 @@ +#!/bin/bash +set -e + +# Detect if a GPU is available via nvidia-smi. +# This works for Linux hosts and Windows (WSL2) with NVIDIA Container Toolkit. +# On Mac (Apple Silicon) or systems without NVIDIA GPUs, this will skip the 'cuda' extra. +if command -v nvidia-smi &> /dev/null && nvidia-smi &> /dev/null; then + echo "GPU detected. Syncing with 'cuda' extra..." + uv sync --frozen --extra cuda +else + echo "No GPU detected or nvidia-smi failed. Syncing without 'cuda' extra..." + uv sync --frozen +fi + +echo "Environment synced successfully." + +# Install pre-commit hooks so they are active in the devcontainer +uv run pre-commit install diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..ebd4eb9 --- /dev/null +++ b/.env.example @@ -0,0 +1,12 @@ +# Brittle Star Project Environment Variables +# Copy this file to .env and fill in your values. +# IMPORTANT: Never commit the actual .env file, it is in .gitignore + +# ---------------------------- # +# Weights and Biases API Key # +# ---------------------------- # +# To find your API key: +# 1. Log in to wandb.ai +# 2. Go to User Settings (https://wandb.ai/settings) +# 3. Scroll down to the "API keys" section +WANDB_API_KEY=your_api_key_here diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..0fc68fc --- /dev/null +++ b/.gitattributes @@ -0,0 +1,14 @@ +# Model weights +*.pt filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.safetensors filter=lfs diff=lfs merge=lfs -text +*.onnx filter=lfs diff=lfs merge=lfs -text +*.h5 filter=lfs diff=lfs merge=lfs -text + +# Media +*.mp4 filter=lfs diff=lfs merge=lfs -text +*.avi filter=lfs diff=lfs merge=lfs -text + +# Datasets +*.csv filter=lfs diff=lfs merge=lfs -text +*.npz filter=lfs diff=lfs merge=lfs -text diff --git a/.github/scripts/prepare_docs.py b/.github/scripts/prepare_docs.py new file mode 100644 index 0000000..2ecefaa --- /dev/null +++ b/.github/scripts/prepare_docs.py @@ -0,0 +1,28 @@ +import os +import glob +import re +import shutil + +folders_to_copy = ["src", "scripts", "configs"] +for folder in folders_to_copy: + if os.path.exists(folder): + shutil.copytree(folder, f"docs/{folder}", dirs_exist_ok=True) + +for filepath in glob.glob("docs/**/*.md", recursive=True): + with open(filepath, "r", encoding="utf-8") as f: + content = f.read() + + # RULE A: Fix links pointing OUT to src/, scripts/, or configs/ + # Logic: Because the folders were moved one level deeper, we remove exactly ONE '../' + content = re.sub( + r"\]\(\.\./((?:\.\./)*)(src|scripts|configs)/([^)]*)\)", r"](\1\2/\3)", content + ) + + # RULE B: Fix links pointing FROM the copied files back TO the original docs/ folder + # Logic: Since these files are now inside docs/, the 'docs/' segment in the path is redundant. + content = re.sub(r"\]\(((?:\.\./)+)docs/([^)]*)\)", r"](\1\2)", content) + + with open(filepath, "w", encoding="utf-8") as f: + f.write(content) + +print("Successfully imported external files and adjusted markdown links.") diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 0000000..2b6b033 --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,23 @@ +name: Lint + +on: + pull_request: + types: [opened, synchronize, ready_for_review] + branches: + - main + - dev + +jobs: + ruff: + if: ${{ !github.event.pull_request.draft }} + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Run Ruff check + uses: astral-sh/ruff-action@v1 + with: + args: "check" + - name: Run Ruff format check + uses: astral-sh/ruff-action@v1 + with: + args: "format --check" diff --git a/.github/workflows/publish-docs.yml b/.github/workflows/publish-docs.yml new file mode 100644 index 0000000..668c977 --- /dev/null +++ b/.github/workflows/publish-docs.yml @@ -0,0 +1,38 @@ +name: Publish docs via GitHub Pages + +on: + push: + branches: + - main + - dev + - docs/* + +jobs: + build: + name: Deploy docs + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + fetch-depth: 0 + ssh-key: ${{ secrets.DEPLOY_KEY }} + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install dependencies + run: pip install mkdocs-material + + - name: Prepare external docs + run: python .github/scripts/prepare_docs.py + + - name: Configure Git identity + run: | + git config --global user.name "github-actions[bot]" + git config --global user.email "github-actions[bot]@users.noreply.github.com" + + - name: Deploy to GitHub Pages + run: mkdocs gh-deploy --force diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..bcf6f2f --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,31 @@ +name: Test + +on: + pull_request: + types: [opened, synchronize, ready_for_review] + branches: + - main + - dev + +jobs: + pytest: + if: ${{ !github.event.pull_request.draft }} + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v5 + with: + version: "latest" + + - name: Setup python + uses: actions/setup-python@v5 + with: + python-version-file: ".python-version" + + - name: Install dependencies + run: uv sync --frozen + + - name: Run tests with pytest + run: uv run pytest tests/ diff --git a/.github/workflows/update_hpc_requirements.yml b/.github/workflows/update_hpc_requirements.yml new file mode 100644 index 0000000..969b433 --- /dev/null +++ b/.github/workflows/update_hpc_requirements.yml @@ -0,0 +1,43 @@ +name: Update HPC requirements + +on: + push: + paths: + - pyproject.toml + branches: + - main + - dev + - "ci/**" + +jobs: + update-hpc-requirements: + runs-on: ubuntu-latest + permissions: + contents: write + pull-requests: write + steps: + - uses: actions/checkout@v4 + with: + ref: ${{ github.head_ref || github.ref_name }} + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install uv + uses: astral-sh/setup-uv@v5 + + - name: Regenerate env/hpc/requirements.txt + run: uv run scripts/hpc/export_requirements.py + + - name: Create Pull Request with updated requirements + uses: peter-evans/create-pull-request@v6 + with: + token: ${{ secrets.GITHUB_TOKEN }} + commit-message: "chore(hpc): update env/hpc/requirements.txt from pyproject.toml" + title: "chore(hpc): update HPC requirements" + body: "Automatically generated pull request to update `env/hpc/requirements.txt` based on recent changes to `pyproject.toml`." + branch: chore/auto-update-hpc-requirements + base: ${{ github.ref_name }} + author: "github-actions[bot] " diff --git a/.gitignore b/.gitignore index da01ac7..d23008f 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,13 @@ +# Model files +artifacts/* +runs/* +wandb/ +outputs/ +multirun/ +metrics/ +adjacency_debug.txt +vids/ + # Python-generated files __pycache__/ *.py[oc] @@ -368,7 +378,6 @@ celerybeat.pid # Environments .env .venv -env/ venv/ ENV/ env.bak/ @@ -468,7 +477,6 @@ tags [Ll]ib [Ll]ib64 [Ll]ocal -[Ss]cripts pyvenv.cfg .venv pip-selfcheck.json @@ -515,4 +523,7 @@ Icon Network Trash Folder Temporary Items .apdisk +*.pdf +# plot directory +poster_plots/ \ No newline at end of file diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..c289bb6 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,21 @@ +repos: + - repo: local + hooks: + - id: ruff-format + name: ruff format (uv) + entry: uv run ruff format + language: system + types: [python] + + - repo: https://github.com/alessandrojcm/commitlint-pre-commit-hook + rev: v9.16.0 + hooks: + - id: commitlint + stages: [commit-msg] + additional_dependencies: ["@commitlint/config-conventional"] + + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: no-commit-to-branch + args: ['--branch', 'main', '--branch', 'dev'] \ No newline at end of file diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 0000000..9b38853 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,7 @@ +{ + "python.testing.pytestArgs": [ + "tests" + ], + "python.testing.unittestEnabled": false, + "python.testing.pytestEnabled": true +} \ No newline at end of file diff --git a/README.md b/README.md index e69de29..6e56334 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,62 @@ +# Brittle Star + +> What is the impact of different levels of controller-modularity on the learning-speed, coordination and tolerance for +defects (e.g. amputations) in brittle-star-like robots trained with Reinforcement Learning? + +## Quick start + +### Local setup + +To set up the UV module, you can run the following command: + +```bash +uv sync --frozen +``` + +## Repository Structure + +```text +. +├── configs/ # Hydra configuration files (YAML) +├── docs/ # Comprehensive documentation and API guides +├── runs/ # Default output directory for Hydra and training artifacts +├── scripts/ # High-level entrypoints for training, simulation, and evaluation +├── src/ +│ ├── brittle_star_project/ # Core library and environment logic +│ │ ├── evaluation/ # Checkpoint evaluation, rollout logic, and metrics persistence +│ │ └── trainers/ # Training loop implementations (e.g., PPO) +│ └── experiment_logger/ # Standalone logging package +└── tests/ # Unit and integration tests +``` + +## Usage + +For detailed instructions on how to use the project, please refer to the **[API Documentation](docs/README.md)**. + +### Quick Start + +1. **Train a model:** + ```bash + uv run python scripts/train.py ppo.learning_rate=0.001 logging.track=true + ``` + +2. **Monitor progress:** + See [Tracking & Monitoring](docs/api/tracking.md). + +3. **Simulate a trained model:** + See [Simulation & Evaluation](docs/api/simulation.md). + +4. **Compare fault tolerance of models:** + See [Checkpoint & Model Evaluation](docs/api/evaluation.md) + +## Results & Reproduction + +See **[docs/api/reproduction.md](docs/api/reproduction.md)** to learn how to access our public [Weights & Biases (WandB) project](https://wandb.ai/SEL3-2026-Groep-4/final-models-v2?nw=96mloffsyq), retrieve specific run parameters, and run the training/evaluation reproduction workflow. + +## HPC + +See **[docs/HPC.md](docs/HPC.md)** for the full guide, including environment setup, cluster selection, interactive debugging, and job submission. + +## Documentation + +Please find all documentation and a starting point for more information in [corresponding README](./docs/README.md). diff --git a/configs/README.md b/configs/README.md new file mode 100644 index 0000000..91208f2 --- /dev/null +++ b/configs/README.md @@ -0,0 +1,54 @@ +# Brittle Star Configuration System + +This project uses **Hydra** for a modular, hierarchical, and strictly-typed configuration system. + +## Core Concepts + +1. **Composition over Inheritance**: Instead of one giant config file, the configuration is composed of small, domain-specific modules (PPO settings, architecture, morphology, etc.). +2. **Strict Typing**: Every configuration is validated against a Python dataclass schema (`ConfigStore`). Misspelled keys throw a `ConfigAttributeError` immediately. +3. **CLI Swapping**: You can swap entire modules or override individual values from the command line without touching code. + +## Directory Structure + +- `main_config.yaml`: The root entry point defining the default composition. +- `experiment/`: High-level experiment settings (seed, device). +- `logging/`: WandB and checkpointing configuration. +- `ppo/`: PPO training hyperparameters. +- `architecture/`: Polymorphic network architectures (centralized vs. decentralized). +- `morphology/`: Physical robot definitions (number of segments, amputations). +- `arena/`: Environment physics and visual settings. +- `environment/`: Task-specific settings (Directed Locomotion, Light Escape). + +## Common Commands + +### Local Debugging +Run a quick test with minimal iterations: +```bash +python scripts/train.py experiment=dev_test ppo=fast +``` + +### Swapping Architectures or Morphologies +Test a decentralized controller on a 3-arm robot: +```bash +python scripts/train.py architecture=decentralized morphology=3_arms +``` + +### HPC Production +Run stable PPO with WandB enabled (HPC submission scripts handle the `hydra.run.dir` redirection): +```bash +python scripts/train.py ppo=stable logging=wandb_enabled +``` + +### Dry-Run Validation +Check if your configuration is valid without starting the simulation: +```bash +python scripts/train.py --cfg job +``` + +## Developer Notes + +- **Adding a new group**: Create a subdirectory in `configs/` and register the new dataclass in `src/brittle_star_project/configs/register_configs.py`. +- **Typo Catching**: If you see a `ConfigAttributeError`, check for typos in your YAML keys or CLI overrides. +- **Output Redirection**: We use `experiment.base_run_dir` to configure where logs and models are stored (defaults to `runs/`). + - To change it locally: `python scripts/train.py experiment.base_run_dir=/path/to/custom/dir` + - On HPC, ensure this points to a fast scratch storage. diff --git a/configs/architecture/centralized.yaml b/configs/architecture/centralized.yaml new file mode 100644 index 0000000..bb4bc0f --- /dev/null +++ b/configs/architecture/centralized.yaml @@ -0,0 +1,22 @@ +# Centralized Actor-Critic Architecture +# Baseline configuration with a single global sensor and motor. + +# Default values are defined in CentralizedConfig dataclass. +# Use this configuration for standard PPO experiments. + +name: "centralized" +sensor: + hidden_dims: [300, 300, 300] + activation: "tanh" + +motor: + hidden_dims: [] + activation: "tanh" + +feature_extractor: + hidden_dims: [300, 300, 300] + activation: "tanh" + +critic: + hidden_dims: [] + activation: "tanh" diff --git a/configs/architecture/decentralized.yaml b/configs/architecture/decentralized.yaml new file mode 100644 index 0000000..70b0fb2 --- /dev/null +++ b/configs/architecture/decentralized.yaml @@ -0,0 +1,32 @@ +# Decentralized Actor Architecture (NerveNet-MLP variant) +# Multi-agent/distributed configuration using local sensors, propagators, and motors. + +# Default values are defined in DecentralizedConfig dataclass. +# Use this configuration for decentralized execution experiments. + +name: "decentralized" +sensor: + hidden_dims: [300, 300, 300] + activation: "tanh" + +propagator: + hidden_dims: [300, 300, 300] + activation: "tanh" + +motor: + hidden_dims: [] + activation: "tanh" + +feature_extractor: + hidden_dims: [300, 300, 300] + activation: "tanh" + +critic: + hidden_dims: [] + activation: "tanh" + +# Synchronous message-passing rounds per control step +message_passing_steps: 4 + +# Connectivity topology (e.g., ring, fully_connected) +topology_type: "fully_connected" diff --git a/configs/arena/default.yaml b/configs/arena/default.yaml new file mode 100644 index 0000000..c42ec3c --- /dev/null +++ b/configs/arena/default.yaml @@ -0,0 +1,8 @@ +# Default Arena Configuration +# Base aquarium environment settings. + +size: [10.0, 5.0] +sand_ground_color: true +attach_target: true +wall_height: 1.5 +wall_thickness: 0.1 diff --git a/configs/centralized-final.yaml b/configs/centralized-final.yaml new file mode 100644 index 0000000..9c50437 --- /dev/null +++ b/configs/centralized-final.yaml @@ -0,0 +1,71 @@ +# Custom Main Configuration +# +# Use with: +# uv run python scripts/train.py --config-name main_config_custom +# +# This keeps the project defaults intact while giving you a single custom +# training entrypoint you can edit freely. + +defaults: + - brittle_star_config + - experiment: base + - logging: default + - evaluation: default + - ppo: default + - architecture: centralized + - morphology: 5_arms_full + - arena: default + - environment: directed_locomotion + - simulation: default + - _self_ + +morphology: + morph_mode: CENTRALIZED + +experiment: + exp_name: "final-models-v2/centralized/" + seed: 42 + torch_deterministic: true + cuda: true + +logging: + track: true + save_model: true + save_checkpoints: true + upload_final_model: true + upload_checkpoints: true + checkpoint_frequency: 10 + wandb_project_name: "final-models-v2" + +evaluation: + evaluate_checkpoints: true + eval_max_steps: 2000 + eval_seed: 0 + +ppo: + learning_rate: 0.0001 + total_timesteps: 16384000 + num_envs: 128 + num_steps: 64 + anneal_lr: true + gamma: 0.99 + gae_lambda: 0.95 + num_minibatches: 32 + update_epochs: 4 + norm_adv: true + clip_coef: 0.2 + clip_vloss: true + ent_coef: 0.001 + vf_coef: 1.0 + max_grad_norm: 0.5 + target_kl: 0.02 + +environment: + simulation_time: 100000.0 + target_distance: 3.0 + +hydra: + job: + chdir: true + run: + dir: ${experiment.base_run_dir}/${experiment.exp_name}/${now:%Y-%m-%d}/${now:%H-%M-%S} diff --git a/configs/environment/dir_loc_further.yaml b/configs/environment/dir_loc_further.yaml new file mode 100644 index 0000000..236f0c5 --- /dev/null +++ b/configs/environment/dir_loc_further.yaml @@ -0,0 +1,2 @@ +simulation_time: 50000.0 +target_distance: 3.0 \ No newline at end of file diff --git a/configs/environment/directed_locomotion.yaml b/configs/environment/directed_locomotion.yaml new file mode 100644 index 0000000..b664ca1 --- /dev/null +++ b/configs/environment/directed_locomotion.yaml @@ -0,0 +1,12 @@ +# Directed Locomotion Environment +# Baseline task setting. + +task: DIRECTED_LOCOMOTION +simulation_time: 100000.0 +num_physics_steps_per_control_step: 10 +time_scale: 2 +camera_ids: [0, 1] +render_size: [480, 640] +joint_randomization_noise_scale: 0.0 +target_distance: 3.0 +light_perlin_noise_scale: 0 diff --git a/configs/environment/light_escape.yaml b/configs/environment/light_escape.yaml new file mode 100644 index 0000000..ade85ba --- /dev/null +++ b/configs/environment/light_escape.yaml @@ -0,0 +1,12 @@ +# Light Escape Environment +# Advanced task requiring movement away from light source. + +task: LIGHT_ESCAPE +simulation_time: 100000.0 +num_physics_steps_per_control_step: 10 +time_scale: 2 +camera_ids: [0, 1] +render_size: [480, 640] +joint_randomization_noise_scale: 0.0 +target_distance: 3.0 +light_perlin_noise_scale: 200 # Must be integer factor of 200 diff --git a/configs/evaluation/default.yaml b/configs/evaluation/default.yaml new file mode 100644 index 0000000..a7c2a00 --- /dev/null +++ b/configs/evaluation/default.yaml @@ -0,0 +1,8 @@ +# Default Evaluation Configuration +# Settings used for checkpoint evaluation during training. + +evaluate_checkpoints: false +# Max number of control steps during evaluation rollout. +eval_max_steps: 2000 +# Seed for deterministic evaluation reset. +eval_seed: 0 diff --git a/configs/evaluation/poster.yaml b/configs/evaluation/poster.yaml new file mode 100644 index 0000000..948e4c7 --- /dev/null +++ b/configs/evaluation/poster.yaml @@ -0,0 +1,25 @@ +# @package evaluation +# Configuration for the models used in the poster comparison. + +# Standard evaluation settings +evaluate_checkpoints: false +eval_max_steps: 5000 +eval_seed: 0 + +# Cross-model comparison settings +# We use 10 episodes to get a more robust average for the final poster results. +comparison_base_seed: 0 +comparison_num_episodes: 10 +comparison_output_csv: "runs/evaluation/comparison.csv" + +# Paths to the .cleanrl_model files to be compared (relative to workspace root). +comparison_models: + - "runs/final-v2-centralized/artifacts/12-19-01_checkpoint_v22/checkpoint_step_230.flax" + - "runs/final-v2-fully-conn/artifacts/14-02-00_checkpoint_v17/checkpoint_step_180.flax" + - "runs/final-v2-ring/artifacts/15-27-03_checkpoint_v21/checkpoint_step_220.flax" + +# Path to the morphologies to evaluate against. +comparison_morphologies: + - "configs/morphology/5_arms_full.yaml" + - "configs/morphology/3_arms.yaml" + - "configs/morphology/2_arms.yaml" diff --git a/configs/experiment/base.yaml b/configs/experiment/base.yaml new file mode 100644 index 0000000..9f31d27 --- /dev/null +++ b/configs/experiment/base.yaml @@ -0,0 +1,7 @@ +# Base Experiment Configuration +# Default values align with ExperimentConfig dataclass. + +exp_name: "brittle_star_ppo" +seed: 1 +torch_deterministic: true +cuda: true diff --git a/configs/experiment/dev_test.yaml b/configs/experiment/dev_test.yaml new file mode 100644 index 0000000..941a1e9 --- /dev/null +++ b/configs/experiment/dev_test.yaml @@ -0,0 +1,7 @@ +# Testing Experiment Configuration +# Quick experiment for local development/testing. + +exp_name: "dev_test_brittle_star" +seed: 42 +torch_deterministic: true +cuda: true diff --git a/configs/experiment/hpc_smoke_test.yaml b/configs/experiment/hpc_smoke_test.yaml new file mode 100644 index 0000000..4264801 --- /dev/null +++ b/configs/experiment/hpc_smoke_test.yaml @@ -0,0 +1,7 @@ +# HPC Smoke Test Configuration +# Uses minimal settings but simulates HPC environment. + +exp_name: "hpc_smoke_test" +seed: 123 +torch_deterministic: true +cuda: true diff --git a/configs/experiment/long_2arm.yaml b/configs/experiment/long_2arm.yaml new file mode 100644 index 0000000..ae4f4b0 --- /dev/null +++ b/configs/experiment/long_2arm.yaml @@ -0,0 +1,6 @@ +# Testing chicken dinner 4 but further distance. + +exp_name: "long2arm" +seed: 123 +torch_deterministic: true +cuda: true \ No newline at end of file diff --git a/configs/fully-connected-final.yaml b/configs/fully-connected-final.yaml new file mode 100644 index 0000000..29983e3 --- /dev/null +++ b/configs/fully-connected-final.yaml @@ -0,0 +1,74 @@ +# Custom Main Configuration +# +# Use with: +# uv run python scripts/train.py --config-name main_config_custom +# +# This keeps the project defaults intact while giving you a single custom +# training entrypoint you can edit freely. + +defaults: + - brittle_star_config + - experiment: base + - logging: default + - evaluation: default + - ppo: default + - architecture: decentralized + - morphology: 5_arms_full + - arena: default + - environment: directed_locomotion + - simulation: default + - _self_ + +architecture: + topology_type: "fully_connected" + +morphology: + morph_mode: FULLY_CONNECTED + +experiment: + exp_name: "final-models-v2/fully-connected/" + seed: 42 + torch_deterministic: true + cuda: true + +logging: + track: true + save_model: true + save_checkpoints: true + upload_final_model: true + upload_checkpoints: true + checkpoint_frequency: 10 + wandb_project_name: "final-models-v2" + +evaluation: + evaluate_checkpoints: true + eval_max_steps: 2000 + eval_seed: 0 + +ppo: + learning_rate: 0.0001 + total_timesteps: 16384000 + num_envs: 128 + num_steps: 64 + anneal_lr: true + gamma: 0.99 + gae_lambda: 0.95 + num_minibatches: 32 + update_epochs: 4 + norm_adv: true + clip_coef: 0.2 + clip_vloss: true + ent_coef: 0.001 + vf_coef: 1.0 + max_grad_norm: 0.5 + target_kl: 0.02 + +environment: + simulation_time: 100000.0 + target_distance: 3.0 + +hydra: + job: + chdir: true + run: + dir: ${experiment.base_run_dir}/${experiment.exp_name}/${now:%Y-%m-%d}/${now:%H-%M-%S} diff --git a/configs/logging/default.yaml b/configs/logging/default.yaml new file mode 100644 index 0000000..2f7de8c --- /dev/null +++ b/configs/logging/default.yaml @@ -0,0 +1,13 @@ +# Default Logging Configuration +# Offline local-only setup (WandB disabled). + +track: false +wandb_project_name: "PPO-Modularity" +wandb_entity: "SEL3-2026-Groep-4" +capture_video: false +save_model: true +save_checkpoints: true +checkpoint_frequency: 100 +upload_final_model: false +upload_checkpoints: false +hf_entity: "" \ No newline at end of file diff --git a/configs/logging/hpc.yaml b/configs/logging/hpc.yaml new file mode 100644 index 0000000..3c03918 --- /dev/null +++ b/configs/logging/hpc.yaml @@ -0,0 +1,9 @@ +track: true +wandb_project_name: "hpc-default" +wandb_entity: "SEL3-2026-Groep-4" +save_model: true +save_checkpoints: true +upload_final_model: true +upload_checkpoints: true +checkpoint_frequency: 100 +hf_entity: "" diff --git a/configs/logging/wandb_enabled.yaml b/configs/logging/wandb_enabled.yaml new file mode 100644 index 0000000..7ddf95f --- /dev/null +++ b/configs/logging/wandb_enabled.yaml @@ -0,0 +1,13 @@ +# WandB Enabled Logging Configuration +# For production/cloud experiments with weights synced. + +track: true +wandb_project_name: "default-project" +wandb_entity: "SEL3-2026-Groep-4" +capture_video: false +save_model: true +save_checkpoints: true +checkpoint_frequency: 100 +upload_final_model: true +upload_checkpoints: false +hf_entity: "" diff --git a/configs/main_config.yaml b/configs/main_config.yaml new file mode 100644 index 0000000..44451f4 --- /dev/null +++ b/configs/main_config.yaml @@ -0,0 +1,23 @@ +# Brittle Star Project - Main Configuration +# This file defines the default composition of the hierarchical configuration. +# Sub-configs are loaded from the relative directories. + +defaults: + - brittle_star_config + - experiment: base + - logging: default + - evaluation: default + - ppo: default + - architecture: centralized + - morphology: 5_arms_full + - arena: default + - environment: directed_locomotion + - obs_bounds: default + - simulation: default + - _self_ + +hydra: + job: + chdir: True + run: + dir: ${experiment.base_run_dir}/${experiment.exp_name}/${now:%Y-%m-%d}/${now:%H-%M-%S} diff --git a/configs/morphology/2_arms.yaml b/configs/morphology/2_arms.yaml new file mode 100644 index 0000000..c301ec4 --- /dev/null +++ b/configs/morphology/2_arms.yaml @@ -0,0 +1,5 @@ +# 2 Arms Morphology Configuration + +segments_per_arm: [4, 0, 4, 0, 0] +use_p_control: true +use_torque_control: false \ No newline at end of file diff --git a/configs/morphology/2_arms_decentralized.yaml b/configs/morphology/2_arms_decentralized.yaml new file mode 100644 index 0000000..9fe3160 --- /dev/null +++ b/configs/morphology/2_arms_decentralized.yaml @@ -0,0 +1,6 @@ +# 2 Arms Morphology Configuration + +segments_per_arm: [4, 0, 4, 0, 0] +use_p_control: true +use_torque_control: false +morph_mode: FULLY_CONNECTED \ No newline at end of file diff --git a/configs/morphology/3_arms.yaml b/configs/morphology/3_arms.yaml new file mode 100644 index 0000000..ebc1665 --- /dev/null +++ b/configs/morphology/3_arms.yaml @@ -0,0 +1,6 @@ +# 3 Arms Morphology Configuration +# Symmetric amputation (arms 1 and 3 removed). + +segments_per_arm: [4, 0, 4, 0, 4] +use_p_control: true +use_torque_control: false diff --git a/configs/morphology/5_arms_damaged.yaml b/configs/morphology/5_arms_damaged.yaml new file mode 100644 index 0000000..2afd643 --- /dev/null +++ b/configs/morphology/5_arms_damaged.yaml @@ -0,0 +1,6 @@ +# 5 Arms Full Morphology Configuration +# Baseline 5-arm brittle star. + +segments_per_arm: [4, 4, 0, 4, 4] +use_p_control: true +use_torque_control: false diff --git a/configs/morphology/5_arms_full.yaml b/configs/morphology/5_arms_full.yaml new file mode 100644 index 0000000..408b7c9 --- /dev/null +++ b/configs/morphology/5_arms_full.yaml @@ -0,0 +1,6 @@ +# 5 Arms Full Morphology Configuration +# Baseline 5-arm brittle star. + +segments_per_arm: [4, 4, 4, 4, 4] +use_p_control: true +use_torque_control: false diff --git a/configs/morphology/5_arms_full_fullconnected.yaml b/configs/morphology/5_arms_full_fullconnected.yaml new file mode 100644 index 0000000..ab08e59 --- /dev/null +++ b/configs/morphology/5_arms_full_fullconnected.yaml @@ -0,0 +1,7 @@ +# 5 Arms Full Morphology Configuration +# Baseline 5-arm brittle star. + +segments_per_arm: [4, 4, 4, 4, 4] +use_p_control: true +use_torque_control: false +morph_mode: FULLY_CONNECTED diff --git a/configs/morphology/partial_amputation.yaml b/configs/morphology/partial_amputation.yaml new file mode 100644 index 0000000..8003035 --- /dev/null +++ b/configs/morphology/partial_amputation.yaml @@ -0,0 +1,6 @@ +# Partial Amputation Configuration +# Random partial amputation for robustness testing. + +segments_per_arm: [4, 2, 4, 4, 4] +use_p_control: true +use_torque_control: false diff --git a/configs/obs_bounds/default.yaml b/configs/obs_bounds/default.yaml new file mode 100644 index 0000000..cd70b0b --- /dev/null +++ b/configs/obs_bounds/default.yaml @@ -0,0 +1 @@ +# Defaults provided by dataclass diff --git a/configs/ppo/chickendinnerwinner.yaml b/configs/ppo/chickendinnerwinner.yaml new file mode 100644 index 0000000..cf5e3d6 --- /dev/null +++ b/configs/ppo/chickendinnerwinner.yaml @@ -0,0 +1,16 @@ +anneal_lr: true +clip_coef: 0.2 +clip_vloss: true +ent_coef: 0.001 +gae_lambda: 0.95 +gamma: 0.99 +learning_rate: 0.0001 +max_grad_norm: 0.5 +norm_adv: true +num_envs: 32 +num_minibatches: 32 +num_steps: 64 +target_kl: 0.02 +total_timesteps: 12288000 +update_epochs: 4 +vf_coef: 1.0 \ No newline at end of file diff --git a/configs/ppo/debug.yaml b/configs/ppo/debug.yaml new file mode 100644 index 0000000..7732fd3 --- /dev/null +++ b/configs/ppo/debug.yaml @@ -0,0 +1,16 @@ +learning_rate: 0.0003 +total_timesteps: 409600 +num_envs: 32 +num_steps: 32 +anneal_lr: true +gamma: 0.99 +gae_lambda: 0.95 +num_minibatches: 32 +update_epochs: 4 +norm_adv: true +clip_coef: 0.2 +clip_vloss: true +ent_coef: 0.005 +vf_coef: 1.0 +max_grad_norm: 0.5 +target_kl: null diff --git a/configs/ppo/default.yaml b/configs/ppo/default.yaml new file mode 100644 index 0000000..50107b4 --- /dev/null +++ b/configs/ppo/default.yaml @@ -0,0 +1,19 @@ +# Default PPO Configuration +# Standard hyperparams from original codebase. + +learning_rate: 0.00025 +total_timesteps: 10000000 +num_envs: 100 +num_steps: 128 +anneal_lr: true +gamma: 0.99 +gae_lambda: 0.95 +num_minibatches: 4 +update_epochs: 4 +norm_adv: true +clip_coef: 0.1 +clip_vloss: true +ent_coef: 0.01 +vf_coef: 0.5 +max_grad_norm: 0.5 +target_kl: null diff --git a/configs/ppo/dev_larger_timesteps_larger_rolloutsteps.yaml b/configs/ppo/dev_larger_timesteps_larger_rolloutsteps.yaml new file mode 100644 index 0000000..b2f629a --- /dev/null +++ b/configs/ppo/dev_larger_timesteps_larger_rolloutsteps.yaml @@ -0,0 +1,16 @@ +learning_rate: 0.0003 +total_timesteps: 1228800 +num_envs: 32 +num_steps: 64 +anneal_lr: true +gamma: 0.99 +gae_lambda: 0.95 +num_minibatches: 32 +update_epochs: 4 +norm_adv: true +clip_coef: 0.2 +clip_vloss: true +ent_coef: 0.005 +vf_coef: 1.0 +max_grad_norm: 0.5 +target_kl: null \ No newline at end of file diff --git a/configs/ppo/smoke_test.yaml b/configs/ppo/smoke_test.yaml new file mode 100644 index 0000000..4b7bb84 --- /dev/null +++ b/configs/ppo/smoke_test.yaml @@ -0,0 +1,19 @@ +# Fast PPO Configuration +# Lower timestep count for quick iterations/testing. + +learning_rate: 0.0005 +total_timesteps: 1024 +num_envs: 32 +num_steps: 32 +anneal_lr: true +gamma: 0.99 +gae_lambda: 0.95 +num_minibatches: 4 +update_epochs: 4 +norm_adv: true +clip_coef: 0.2 +clip_vloss: true +ent_coef: 0.01 +vf_coef: 0.5 +max_grad_norm: 0.5 +target_kl: null diff --git a/configs/ppo/stable.yaml b/configs/ppo/stable.yaml new file mode 100644 index 0000000..35cade4 --- /dev/null +++ b/configs/ppo/stable.yaml @@ -0,0 +1,19 @@ +# Stable PPO Configuration +# Standard hyperparams with lower LR and larger batch. + +learning_rate: 0.0001 +total_timesteps: 10000000 +num_envs: 100 +num_steps: 256 +anneal_lr: true +gamma: 0.99 +gae_lambda: 0.95 +num_minibatches: 8 +update_epochs: 4 +norm_adv: true +clip_coef: 0.1 +clip_vloss: true +ent_coef: 0.01 +vf_coef: 0.5 +max_grad_norm: 0.5 +target_kl: null diff --git a/configs/ring-final.yaml b/configs/ring-final.yaml new file mode 100644 index 0000000..a0d852a --- /dev/null +++ b/configs/ring-final.yaml @@ -0,0 +1,74 @@ +# Custom Main Configuration +# +# Use with: +# uv run python scripts/train.py --config-name main_config_custom +# +# This keeps the project defaults intact while giving you a single custom +# training entrypoint you can edit freely. + +defaults: + - brittle_star_config + - experiment: base + - logging: default + - evaluation: default + - ppo: default + - architecture: decentralized + - morphology: 5_arms_full + - arena: default + - environment: directed_locomotion + - simulation: default + - _self_ + +architecture: + topology_type: "ring" + +morphology: + morph_mode: RING + +experiment: + exp_name: "final-models-v2/ring/" + seed: 42 + torch_deterministic: true + cuda: true + +logging: + track: true + save_model: true + save_checkpoints: true + upload_final_model: true + upload_checkpoints: true + checkpoint_frequency: 10 + wandb_project_name: "final-models-v2" + +evaluation: + evaluate_checkpoints: true + eval_max_steps: 2000 + eval_seed: 0 + +ppo: + learning_rate: 0.0001 + total_timesteps: 16384000 + num_envs: 128 + num_steps: 64 + anneal_lr: true + gamma: 0.99 + gae_lambda: 0.95 + num_minibatches: 32 + update_epochs: 4 + norm_adv: true + clip_coef: 0.2 + clip_vloss: true + ent_coef: 0.001 + vf_coef: 1.0 + max_grad_norm: 0.5 + target_kl: 0.02 + +environment: + simulation_time: 100000.0 + target_distance: 3.0 + +hydra: + job: + chdir: true + run: + dir: ${experiment.base_run_dir}/${experiment.exp_name}/${now:%Y-%m-%d}/${now:%H-%M-%S} diff --git a/configs/simulation/default.yaml b/configs/simulation/default.yaml new file mode 100644 index 0000000..1ebc263 --- /dev/null +++ b/configs/simulation/default.yaml @@ -0,0 +1,30 @@ +# Default Simulation Settings +# These values are used by scripts/simulate.py + +# Path to the trained model (optional) +model_path: null + +# Script behavior +headless: false +# In headless mode this is required; in viewer mode null means "infinite". +max_steps: null + +# Optional: override morphology for amputation experiments. +# Points to a morphology config YAML file (e.g. configs/morphology/3_arms.yaml). +# If null, the training morphology from the model's metadata is used. +morphology_override: null + +# Video recording (requires [evaluation] extra) +record_video: false +# When null, video is saved in a per-model evaluation folder alongside the model. +video_output_path: null +# Camera ID to use for video recording (1 is usually the close-up camera) +camera_id: 1 + +video_width: 640 +video_height: 80 +video_fps: 60 + +# Optional override for the metadata YAML file path. +# If null, the script looks for `_metadata.yaml` alongside the model_path. +metadata_path: null diff --git a/docs/.gitkeep b/docs/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/docs/CONTRIBUTING.md b/docs/CONTRIBUTING.md new file mode 100644 index 0000000..5903790 --- /dev/null +++ b/docs/CONTRIBUTING.md @@ -0,0 +1,44 @@ +# Contribution Guidelines + +This document outlines the contribution protocols for the scientific software engineering project focusing on bio-inspired control architectures for brittle-star-like robots. The primary objective of this project is to produce scientific insight, rather than a commercial product. + +## 1. Scientific Context & Methodology + +* **Research Focus:** The goal is to study how controller modularity affects learning speed, coordination, and fault tolerance in brittle-star locomotion. +* **Hypothesis-Driven Design:** Clear hypotheses must dictate a structured methodology and rigorous evaluation. All design decisions must be formally documented prior to implementation. +* **Scaffolding Approach:** Development must start with simple setups before progressively increasing the complexity of environments and morphologies. +* **Evaluation of Results:** Negative results possess scientific validity when thoroughly analyzed. If a controller fails to learn locomotion, providing a comprehensive analysis of the failure is considered a strong scientific contribution. +* **Reproducibility:** Contributors must utilize fixed library versions. Configuration systems (such as json, gin, or yaml) must be employed to ensure reproducible runs. + +## 2. Clean Code & Code Quality + +Code readability is paramount, as code is read far more frequently than it is written. + +* **Naming Conventions:** Variables and functions must utilize consistent, intention-revealing names. A long, descriptive name is strictly preferred over a short name accompanied by a comment. +* **Function Design:** Functions must be modular and adhere to the single responsibility principle. Arguments must be minimized, and boolean flag arguments controlling behavior should be avoided. +* **Commenting:** Code must document the "how," while comments are strictly reserved for documenting the "why". Commented-out code is prohibited and must be deleted via version control. +* **YAGNI:** Contributors must adhere to the "You Aren't Gonna Need It" (YAGNI) principle and actively avoid premature optimization. +* **Notebooks:** Jupyter Notebooks are strictly limited to quick prototyping, tutorials, demonstrations, or post-processing analysis. They are explicitly forbidden for general software development because they discourage modularity. + +## 3. Version Control & Repository Structure + +* **Git Practices:** Commits must be frequent and small. Each commit should relate to exactly one piece of functionality. +* **Branching Strategy:** The `dev` branch serves as the integration branch for pushing and merging code. Only stable releases may be pushed to the `main` branch. +* **Artifact Management:** Data files, trained models, and large datasets must never be committed directly to Git. Git Large File Storage (LFS) must be used for tracking large files. **All developers must have `git-lfs` installed locally** (see [DEVELOPMENT.md](./DEVELOPMENT.md) for setup). +* **Repository Layout:** The repository must maintain the following core directories: `src/` for algorithms, `env/` for MuJoCo wrappers, `config/` for experiment configurations, `experiments/` for scripts, `docs/` for Doxygen or ReadTheDocs documentation, and `tests/` for unit tests. + +## 4. Architecture & Tooling + +* **Algorithms & Frameworks:** Proximal Policy Optimization (PPO) is the recommended baseline algorithm. CleanRL should be used as a starting point and adapted for continuous action spaces. All Artificial Neural Network (ANN) controller architectures must be implemented using Flax. +* **Simulation:** The simulation environment utilizes a MuJoCo brittle star. XML MuJoCo structures must remain realistic and respect morphological constraints. +* **Experiment Tracking:** Weights & Biases (wandb) must be utilized for tracking and logging all experiments. +* **Code Styling:** All code must conform to the chosen style guide (Google standard). This is enforced via `uv` using **ruff** and pre-commit hooks. + +## 5. AI-Assisted Development & Code Review + +This project supports AI-assisted development to enhance productivity, but contributors must take full responsibility for all AI-generated outputs. + +* **Self-Review Requirement:** Contributors must thoroughly self-review all AI-assisted code, documentation, and configurations before requesting peer review. This includes verifying correctness, adherence to project standards, scientific validity, and integration with existing code. +* **Quality Standards:** AI-generated content must meet the same rigorous standards as manually written code, including proper testing, documentation, and alignment with the scientific methodology outlined in Section 1. +* **Available Skills:** This project provides specific AI skills for common tasks (located in `.agents/skills/`), including linting and testing workflows. Contributors should leverage these skills to maintain consistency and quality. +* **Transparency:** When using AI assistance for complex algorithmic decisions or scientific design choices, contributors should document the rationale in commit messages or code comments where appropriate. diff --git a/docs/DEVELOPMENT.md b/docs/DEVELOPMENT.md new file mode 100644 index 0000000..a6ba2a9 --- /dev/null +++ b/docs/DEVELOPMENT.md @@ -0,0 +1,72 @@ +# Development Guide + +This guide outlines how to set up the development environment for this project, prioritizing **reproducible builds**, **environment parity**, and **cross-hardware compatibility**. + +## Reproducibility & uv + +This project uses [uv](https://github.com/astral-sh/uv) to manage dependencies and virtual environments. The `uv.lock` file is the absolute source of truth for package versions and must always be committed. + +### Source of Truth + +- **Never modify `uv.lock` manually.** +- To add a dependency, run `uv add `. +- To update dependencies, run `uv lock --upgrade`. +- To sync your environment with the lockfile, run `uv sync --frozen`. + +## Git LFS (Critical) + +**All developers must have Git LFS installed locally.** This repository tracks model weights (`.pt`, `.safetensors`, etc.), recordings (`.mp4`), and datasets using Git LFS. + +- **Setup:** Run `git lfs install` after cloning this repository. If you are using the `.devcontainer` or `flake.nix`, LFS is typically available automatically. +- If you clone without LFS installed, run `git lfs pull` after installation to fetch the actual data files instead of the small pointer files. + +## Devcontainer Setup (Recommended) + +The devcontainer provides an identical experience to local development but with all system dependencies pre-configured. It automatically detects your hardware (GPU vs CPU) and syncs the appropriate dependencies. + +### Prerequisites + +- Docker Desktop or Docker Engine. +- [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) (for GPU support). + +### Setup for VS Code + +1. Install the [Dev Containers](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers) extension. +2. Open the project and click **Reopen in Container**. +3. On first launch, the `post-create.sh` script will: + - Detect if an NVIDIA GPU is available via `nvidia-smi`. + - Run `uv sync --frozen --extra cuda` if a GPU is found. + - Run `uv sync --frozen` otherwise. +4. The environment is stored in a **named volume** for `.venv` to ensure persistence and performance. + +### Setup for JetBrains IDEs + +1. The IDE will detect the `.devcontainer/devcontainer.json` file. +2. The environment is pre-configured to point to `/workspaces/project/.venv`. +3. The hardware-aware sync will run automatically during container creation. + +## Local Development (Alternative) + +If you prefer not to use Docker: + +1. Install [uv](https://docs.astral.sh/uv/getting-started/installation/). +2. Run `uv sync --frozen` (CPU) or `uv sync --frozen --extra cuda` (GPU). + +## Hardware Acceleration (JAX) + +Verify your setup by running the JAX initialization test: + +```bash +uv run pytest tests/test_jax_init.py +``` + +In the devcontainer, this will succeed on both CPU and GPU. A `GpuDevice` is expected if a GPU is detected and the `cuda` extra was installed. + +## Logging & Monitoring + +This project uses a unified logging system through the `experiment_logger` package. + +- **Usage in Code**: To use the logger in your scripts, refer to the [package README](../src/experiment_logger/README.md) for the API reference. +- **WandB/TensorBoard Setup**: For information on how to configure tracking for experiments, see the [Tracking & Monitoring API Guide](./api/tracking.md). + +The logger automatically detects if it is running in an interactive terminal or a non-interactive environment (like an HPC Slurm job), adjusting progress bars and fallback modes accordingly. diff --git a/docs/HPC.md b/docs/HPC.md new file mode 100644 index 0000000..e7a3372 --- /dev/null +++ b/docs/HPC.md @@ -0,0 +1,89 @@ +# HPC Guide + +Full documentation: [https://docs.hpc.ugent.be/](https://docs.hpc.ugent.be/) + +## Storage Overview + +- **Run Outputs**: Written to `$VSC_SCRATCH` during the job (fast I/O) and copied to `$VSC_DATA` at the end for persistence. +- **Virtual Environments**: Managed on **`$VSC_DATA`** by mirroring configuration files. This avoids the 3GB home quota without requiring symlinks in the project root. + +## Initial Environment Setup + +Run **once** after cloning the repository. This script handles all modules, mirroring, and environment synchronization. + +```bash +# Option A: Interactive (on a compute node) +module swap cluster/donphan # Debug cluster (CPU only) +# OR for GPU clusters: +# module swap cluster/joltik +# module swap cluster/accelgor +# module swap cluster/litleo + +qsub -I -l nodes=1:gpus=1 # Only for GPU clusters +cd "${PBS_O_WORKDIR}" +bash scripts/hpc/install.sh + +# Option B: Batch (Run in background) +# NOTE: GPU clusters (joltik/accelgor/litleo) require -l gpus=1 at runtime +qsub -l gpus=1 scripts/hpc/install.sh +``` + +## Production vs. Debug Clusters + +Our scripts are cluster-agnostic and do **not** have hardcoded GPU requirements. Instead, you must request GPUs at runtime using the `-l gpus=1` flag when submitting to a production GPU cluster. + +### Debugging (Donphan) + +The `donphan` cluster does not support GPUs. Simply run the scripts without extra resource flags: +```bash +module swap cluster/donphan +qsub scripts/hpc/train.pbs +``` + +### Production (Joltik, Accelgor, Litleo) + +These clusters provide GPU acceleration and **require** a GPU request at runtime: +```bash +module swap cluster/joltik # or accelgor/litleo +qsub -l gpus=1 scripts/hpc/train.pbs +``` + +## Interactive Debugging + +To activate your environment for interactive work, simply run the same `install.sh` script. + +```bash +qsub -I -l nodes=1:ppn=4 -l walltime=1:00:00 +cd "$PBS_O_WORKDIR" +bash scripts/hpc/install.sh +``` + +### Verification Commands + +After installation, run these commands to ensure your environment is set up correctly: + +1. **Verify Quota Safety**: + ```bash + ls -d venvs 2>/dev/null && echo "FAIL" || echo ">>> PASS: Project root is clean." + ``` + +2. **Verify Library Versions (NumPy Fix)**: + ```bash + python -c "import numpy; print(f'NumPy: {numpy.__version__}')" + # Expected: 2.x.x (Venv version), not 1.2x (System version) + ``` + +3. **Verify GPU Access**: + ```bash + python -c "import torch, jax; print(f'GPU: {torch.cuda.is_available()}'); print(f'JAX: {jax.devices()}')" + ``` + +## Managing Dependencies + +`env/hpc/requirements.txt` is auto-generated from `pyproject.toml`. To regenerate: + +```bash +uv run scripts/hpc/export_requirements.py +``` + +Modules listed in `env/hpc/modules.txt` are automatically excluded from the pip requirements to save space and use HPC-optimized binaries. diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 0000000..c2edd2f --- /dev/null +++ b/docs/README.md @@ -0,0 +1,52 @@ +# Documentation + +Welcome to the Brittle Star project documentation. This codebase contains the implementations and research for the scientific evaluation of controller modularity in brittle-star-like robots trained using Reinforcement Learning. + +For the core codebase, scripts, and contribution history, visit our [GitHub Repository](https://github.com/SELab-3-2026/SEL3-2026-Groep-4). + +## Core Requirements & Guides + +- **[Installation Instructions](./DEVELOPMENT.md)**: Steps to set up your development environment locally or in a devcontainer using `uv`, including GPU configuration. For High-Performance Computing (HPC) setup details, see the **[HPC Guide](./HPC.md)**. +- **[How to Run Experiments](./api/training.md)**: A complete guide on running training jobs, setting custom hyperparameters, and overriding config options using Hydra. +- **[Results & Reproduction](./api/reproduction.md)**: Guide on how to access our public WandB training runs table and reproduce our training and evaluation phases (determining the best checkpoint vs. comparing architectures). +- **[Contribution Guidelines](./CONTRIBUTING.md)**: Standards, rules, and best practices for developing and adding code to the repository. +- **[Repository Structure](#repository-structure)**: Overview of the directories and files within the codebase. + +## Repository Structure + +```text +. +├── configs/ # Hydra configuration files (YAML) +├── docs/ # Comprehensive documentation and API guides +├── runs/ # Default output directory for Hydra and training artifacts +├── scripts/ # High-level entrypoints for training, simulation, and evaluation +├── src/ +│ ├── brittle_star_project/ # Core library and environment logic +│ │ ├── evaluation/ # Checkpoint evaluation, rollout logic, and metrics persistence +│ │ └── trainers/ # Training loop implementations (e.g., PPO) +│ └── experiment_logger/ # Standalone logging package +└── tests/ # Unit and integration tests +``` + +## Design & architecture (`/design`) + +If you are interested in the "why did you do it like this?" + +- [Actor-Critic Architecture](./design/actor-critic.md): Description of the actor-critic pipeline. +- [Communication Scheme](./design/communication.md): Message propagation, Nerve-Net style. +- [Modularity & Topology](./design/controllers.md): Macroscopic brain topology, centralized, arm-level, segment-level. +- [Input & Action Spaces](./design/input_action_spaces.md): Description of the model's input and output. +- [Reinforcement Learning Algorithm](./design/learning_algorithm.md): RL techniques, i.e. PPO. +- [Reward Function & Observation Space](./design/reward_function.md): Goals, fitness tracking, and reward structures. + +## API reference (`/api`) + +If you are interested in the "how do I use it?" + +- [Brittle Star Environment](./api/environment.md): MuJoCo environment interaction and configuration. +- [Training Models](./api/training.md): How to configure and run experiments. +- [Tracking & Monitoring](./api/tracking.md): Setting up WandB and TensorBoard to monitor runs. +- [Checkpoint & Model Evaluation](./api/evaluation.md): Evaluating checkpoints and comparing fault tolerance. +- [Interactive Simulation & Visualization](./api/simulation.md): Visualizing models in the MuJoCo viewer or rendering simulation videos. +- [Analysis & Plotting Tools](./api/analysis.md): Comparing checkpoints and generating plots. +- [Results & Reproduction](./api/reproduction.md): Accessing WandB results and running reproduction pipelines. \ No newline at end of file diff --git a/docs/api/analysis.md b/docs/api/analysis.md new file mode 100644 index 0000000..9c88895 --- /dev/null +++ b/docs/api/analysis.md @@ -0,0 +1,89 @@ +# Analysis & Plotting Tools + +This guide outlines the tools available for analyzing experimental data and generating poster-quality visualizations for the Brittle Star project. + +## Shared Configuration + +All plotting scripts share a central configuration in `scripts/plots/plot_config.py`. This file defines: + +- **Color Palette:** A color-blind friendly, high-contrast palette for different architectures. +- **Typography:** Consistent font sizes and styles tailored for A0 posters. +- **Markers:** Shared visual indicators, such as the ★ used for best performers. + +## Comparison Visualization + +The `scripts/plots/analyze_comparisons.py` script generates grouped bar charts comparing the performance of different architectures across various morphologies. + +### Usage + +Run the script from the root of the project, providing the path to your evaluation CSV: + +```bash +# Basic usage (saves PNG and SVG to runs/evaluation/plots/) +uv run python scripts/plots/analyze_comparisons.py path/to/results.csv + +# Advanced usage for Figma/Poster integration +uv run python scripts/plots/analyze_comparisons.py path/to/results.csv \ + --output_dir docs/assets/plots/ \ + --font_size 30 \ + --fig_width 14 \ + --fig_height 10 +``` + +### CLI Arguments + +- `input_csv`: (Required) Path to the CSV file containing evaluation results. +- `--output_dir`, `-o`: Directory where plots will be saved (default: `runs/evaluation/plots`). +- `--show_titles`: Include titles in the plots. Default is **False**, as titles are typically added natively in design tools like Figma. +- `--font_size`: Base font size in points (default: 28). +- `--fig_width` / `--fig_height`: Physical dimensions of the plot in inches. Match these to your Figma layout to maintain exact font sizes. + +### Outputs + +The script generates four key plots, each saved as both `.png` and `.svg`: + +1. **Forward Velocity:** Grouped bar chart (cm/s). +2. **Accumulated Reward:** Mean cumulative reward. +3. **Success Rate:** Target acquisition percentage. +4. **Distance Remaining:** Navigational accuracy. + +--- + +## Convergence Analysis + +The `scripts/plots/analyze_convergence.py` script determines the convergence point of training runs. + +### Usage + +```bash +uv run python scripts/plots/analyze_convergence.py --output_dir runs/convergence/ +``` + +### Configuration + +- **File Mapping:** The script uses hardcoded paths in the `FILE_MAPPING` dictionary. Update these paths to point to your specific run evaluation files. +- **CLI Arguments:** Supports the same `--show_titles`, `--font_size`, and `--fig_width/height` flags as the comparison script. + +### Outputs + +Generates three plots (PNG & SVG): + +1. `convergence_comparison`: Grouped horizontal bar chart. +2. `progress_reward_curves`: Line plots of reward over time. +3. `progress_velocity_curves`: Line plots of velocity over time. + +--- + +## Poster Integration (Figma) + +### SVG & Scaling + +We recommend using the **SVG** outputs for poster design in Figma: + +1. **No Resolution Loss:** SVGs are vector-based and will remain sharp at any size. +2. **Native Text:** Text in the SVG imports as native text layers in Figma. +3. **Exact Font Matching:** To ensure a `28pt` font in the plot matches a `28pt` font in your poster, set the `--fig_width` and `--fig_height` to match the physical dimensions of the plot box in your Figma layout. +4. **Editable:** You can "Ungroup" the SVG in Figma to manually move labels, adjust colors, or tweak individual bars. + +### Image Placeholders +The comparison charts include light-gray square placeholders below the X-axis. These are designed as guides; in Figma, you can drop your morphology renders or illustrations directly on top of these squares. diff --git a/docs/api/environment.md b/docs/api/environment.md new file mode 100644 index 0000000..e54a698 --- /dev/null +++ b/docs/api/environment.md @@ -0,0 +1,30 @@ +# Brittle star environment + +## Creation + +The environment package contains a factory class `BrittleStarEnvFactory` +that creates instances of the environment/morphologies/... It uses the +configuration classes defined in `env_config.py` to create the instances. + +## Configuration + +The data classes in `env_config` have default values as stated in the tutorials. + +* MorphologyConfig: configuration for the morphology of the brittle star. Contains +number of arms, number of segments per arm, and control mode. +* ArenaConfig: configuration for the arena. Sets the size of the arena, whether to +set the ground floor to sand, attach a target and sizes of the walls. +* EnvConfig: configuration for the environment. These set shared settings +such as camera locations, simulation time and the task. + +## Backend and Task enums + +The Backend enum specifies either an MJC or MJX backend. + +* MJC: runs on CPU +* MJX: uses jax on the gpu + +The Task enum specifies which task to use. 2 items are present: + +* DIRECTED_LOCOMOTION: move to a target location +* LIGHT_ESCAPE: situation where the robot must move to a darker location diff --git a/docs/api/evaluation.md b/docs/api/evaluation.md new file mode 100644 index 0000000..1828a7f --- /dev/null +++ b/docs/api/evaluation.md @@ -0,0 +1,60 @@ +# Checkpoint & Model Evaluation + +This guide covers how to evaluate trained brittle star models, with a focus on measuring defect tolerance (amputations) across different controller architectures. + +## Checkpoint Evaluation (During Training) + +The `PPOTrainer` can automatically evaluate every saved checkpoint using the fast MJX backend. This is enabled via configuration. + +### Configuration + +In your experiment config or via CLI: +```bash +python scripts/train.py evaluation.evaluate_checkpoints=true evaluation.eval_max_steps=5000 +``` + +Results are saved to `runs//metrics/checkpoint_evaluation.csv` and synced to Weights & Biases if enabled. + +## Cross-Model & Fault Tolerance Analysis + +To measure how well different controllers handle damage (amputations), use `scripts/compare_models.py`. This script performs a grid search over models x morphologies. + +1. Create or update a YAML file in `configs/evaluation`. +2. Run the benchmark: + +```bash +python scripts/compare_models.py evaluation=poster +``` + +The script will evaluate every combination of model and morphology for the specified number of episodes. + +The results are saved to a CSV (default: `metrics/model_comparison.csv`). + +### CSV Schema + +| Column | Description | +|-----------------------|--------------------------------------------------------------| +| `model_path` | Path to the trained weights. | +| `architecture` | The `morph_mode` of the model (e.g., `CENTRALIZED`, `RING`). | +| `arm_0` ... `arm_4` | Number of segments in each arm slot (0 = amputated). | +| `num_active_arms` | Total number of arms with segments > 0. | +| `seed` | The episode seed. | +| `eval_return` | Accumulated shaped reward. | +| `approx_max_velocity` | Average velocity: `(initial_dist - final_dist) / steps`. | +| `reached_target` | Whether the robot finished within the success radius. | + +## Post-hoc Checkpoint Scanning + +If you need to re-evaluate every saved checkpoint in a run (e.g., to generate a learning curve with different metrics): + +```bash +python scripts/evaluate_checkpoints.py \ + simulation.model_path=runs//final_model.flax \ + evaluation.eval_max_steps=2000 +``` + +This script scans the `checkpoints/` directory of the specified run and evaluates every `.flax` file it finds using the model's training morphology. + +--- + +For a step-by-step walkthrough on using these evaluation phases to reproduce our project results, see the **[Results & Reproduction Guide](./reproduction.md)**. diff --git a/docs/api/reproduction.md b/docs/api/reproduction.md new file mode 100644 index 0000000..87df150 --- /dev/null +++ b/docs/api/reproduction.md @@ -0,0 +1,108 @@ +# Results & Reproduction + +This guide explains how to access our official training logs and reproduce our results. + +Our official training runs, model configurations, and metrics are publicly hosted on Weights & Biases (WandB). + +--- + +## Weights & Biases (WandB) Project + +All experiments, final models, and training logs are tracked in our public WandB project: + +* **Official Runs Table**: [WandB final-models-v2 Table](https://wandb.ai/SEL3-2026-Groep-4/final-models-v2/table?nw=96mloffsyq) + +This page lists the verified runs with their architecture types, morphology definitions, evaluation metrics, and final model performance. + +### How to Reproduce a Run from WandB + +Weights & Biases provides a built-in feature to extract the exact parameters and commands used for any given run: + +1. Open the [WandB final-models-v2 Table](https://wandb.ai/SEL3-2026-Groep-4/final-models-v2/table?nw=96mloffsyq). +2. Click on the name of the run you wish to reproduce to open its detail page. +3. In the top-right corner of the run header (next to the run name, not the main workspace header), click the **three dots (`...`)** menu. +4. Select **"Reproduce run"**. This will display the exact command-line arguments and configuration settings used to execute that run. + +--- + +## Local & HPC Reproduction Workflow + +To reproduce our training and evaluation phases locally or on an HPC cluster, follow the procedures below. + +### 1. Environment Setup + +To ensure identical package versions (including JAX, Flax, and MuJoCo), sync your environment using the lockfile: + +```bash +uv sync --frozen +``` + +### 2. Training Phase + +Run the training script using the exact parameters retrieved from WandB's "Reproduce run" page or from a downloaded `_metadata.yaml` file: + +```bash +uv run python scripts/train.py experiment=my_experiment ppo.learning_rate=0.001 experiment.seed=42 +``` + +--- + +## Evaluation Phases + +Reproducing our evaluation results is divided into two distinct phases: + +### Phase 1: Determining the Best Checkpoint + +During training, checkpoints are saved at regular intervals. To determine which of these checkpoints performed the best: + +1. **Evaluate Checkpoints Post-Training**: + If checkpoint evaluation was not run during training, scan the completed run's checkpoints folder by pointing to the final model path: + + ```bash + uv run python scripts/evaluate_checkpoints.py simulation.model_path=runs/your_run_dir/final_model.flax + ``` + + This script runs deterministic rollouts for every checkpoint in `runs/your_run_dir/checkpoints/`. + +2. **Locate the Results**: + The evaluations are saved to: + + ```text + runs/your_run_dir/metrics/checkpoint_evaluation.csv + ``` + + Analyze this CSV to find the checkpoint iteration with the highest average return or target success rate. This checkpoint will be used for cross-architecture comparisons. + +### Phase 2: Comparing Checkpoints Between Architectures + +Once the best checkpoints for each architecture are identified, they are compared under shared, standardized environments (including fault tolerance checks such as leg amputations). + +1. **Configure the Comparison Models**: + Open or create an evaluation config file (e.g., `configs/evaluation/poster.yaml`) and add the paths to the best checkpoints: + + ```yaml + # configs/evaluation/poster.yaml + evaluation: + comparison_models: + - runs/run_arch_centralized/checkpoints/checkpoint_best.flax + - runs/run_arch_decentralized/checkpoints/checkpoint_best.flax + ``` + +2. **Execute the Comparison Script**: + Run the comparison script using your config: + + ```bash + uv run python scripts/compare_models.py evaluation=poster + ``` + + This script runs multiple sequential evaluation episodes (defined by `comparison_num_episodes` starting at `comparison_base_seed`) for every model across the selected morphologies. + +3. **Analyze Comparison Metrics**: + The script writes a consolidated CSV file to `metrics/model_comparison.csv` containing: + + * **`eval_return`**: The cumulative return. + * **`approx_max_velocity`**: The distance covered per step. + * **`reached_target`**: Navigational success rates. + * **`arm_0` to `arm_4`**: Active segments per arm (indicating damage/amputations). + +This CSV can then be passed to the plotting scripts (e.g., `scripts/plots/analyze_comparisons.py`) to generate visualization plots. For details on configuration and outputs, see the **[Analysis & Plotting Guide](./analysis.md)**. diff --git a/docs/api/simulation.md b/docs/api/simulation.md new file mode 100644 index 0000000..cc87023 --- /dev/null +++ b/docs/api/simulation.md @@ -0,0 +1,54 @@ +# Interactive Simulation & Visualization + +The simulation pipeline allows you to visualize trained models and observe their behavior under various conditions. + +## Overview + +The simulation pipeline is metadata-driven. Training-specific configurations (morphology, arena, environment, etc.) are automatically loaded from the `_metadata.yaml` file associated with the model checkpoint. + +## Basic Simulation + +To simulate a model in the MuJoCo viewer: + +```bash +uv run scripts/simulate.py simulation.model_path=runs/your_run/final_model.flax +``` + +## Amputation & Morphology Overrides + +You can test trained models on different morphologies (e.g., amputating legs) by providing a morphology override. The observations will be automatically padded up to the training morphology's dimensions: + +```bash +uv run scripts/simulate.py \ + simulation.model_path=runs/your_run/final_model.flax \ + simulation.morphology_override=configs/morphology/3_arms.yaml +``` + +## Video Recording + +Recording videos requires the `[evaluation]` extra: + +```bash +uv run scripts/simulate.py \ + simulation.model_path=runs/your_run/final_model.flax \ + simulation.record_video=true \ + simulation.max_steps=1000 +``` + +Videos and evaluation metadata are stored in timestamped folders alongside the model: +`runs/your_run/final_model_evaluations/eval_/simulation.mp4` + +### Top-Down and Follow Cameras + +Using the following script, you can render a top-down and follow camera view for multiple models at once: + +```bash +uv run scripts/poster_visualisations/render_poster_videos.py \ + runs/final-models/centralized/.../final_model.flax \ + runs/final-models/fully-connected/.../final_model.flax \ + runs/final-models/ring/.../final_model.flax \ + --max-steps 10000 --width 640 --height 480 --fps 60 \ + --output-root vids/poster/ +``` +For batch evaluation, checkpoint analysis, and cross-model architecture comparisons, see the **[Checkpoint & Model Evaluation Guide](./evaluation.md)**. + diff --git a/docs/api/tracking.md b/docs/api/tracking.md new file mode 100644 index 0000000..77178f7 --- /dev/null +++ b/docs/api/tracking.md @@ -0,0 +1,64 @@ +# Tracking & Monitoring + +This guide explains how to monitor your experiments using Weights & Biases (WandB) and TensorBoard. + +## Weights & Biases (WandB) + +WandB is used for online synchronization and visualization of training metrics. + +### Authorization + +Export your API key in your terminal to enable WandB synchronization: + +```bash +export WANDB_API_KEY=your_copied_api_key_here +``` + +Alternatively, you can log in using the CLI: + +```bash +uv run wandb login +``` + +### Enabling Tracking + +To enable online sync during a training run, set `logging.track=true` on the command line: + +```bash +uv run python scripts/train.py logging.track=true +``` + +You can also configure your project and entity: + +```bash +uv run python scripts/train.py \ + logging.track=true \ + logging.wandb_project_name="MyProject" \ + logging.wandb_entity="my-team" +``` + +These can also be set in your configuration YAML file under the `logging` key. + +## Local Monitoring with TensorBoard + +All runs are recorded locally in the `runs/` directory (or the directory specified in `experiment.base_run_dir`). You can view scalars and other metrics with TensorBoard: + +```bash +tensorboard --logdir runs/ +``` + +Access the interface at [http://localhost:6006](http://localhost:6006). + +### CLI Exploration Tool + +For quick diagnostics or to export data to CSV without launching the full TensorBoard UI, you can use the `explore_tensorboard.py` script: + +```bash +uv run python scripts/analysis/explore_tensorboard.py runs/your_run_name/ +``` + +See the detailed description in [`/scripts/analysis/README.md`](../../scripts/analysis/README.md). + +## Developer Logging API + +For details on the developer API of our internal logging library (how backend routing, checkpoint synchronization, and singleton initialization works), see the **[Experiment Logger API Guide](../../src/experiment_logger/README.md)**. diff --git a/docs/api/training.md b/docs/api/training.md new file mode 100644 index 0000000..b93b910 --- /dev/null +++ b/docs/api/training.md @@ -0,0 +1,64 @@ +# Training Models + +This guide covers how to configure and run training experiments for the Brittle Star project using Hydra-based configurations. + +## Configuration + +The project uses a modular configuration system powered by [Hydra](https://hydra.cc/). Instead of passing many command-line flags, you select and override configuration groups. + +For a detailed guide on the structure, validation, and usage of our Hydra configuration files, see the **[Brittle Star Configuration System Guide](../../configs/README.md)**. + +### Creating a Custom Experiment + +1. **Create a new experiment file:** + Create a file at `configs/experiment/my_experiment.yaml`. You can copy an existing one as a template: + ```bash + cp configs/experiment/base.yaml configs/experiment/my_experiment.yaml + ``` + +2. **Edit `configs/experiment/my_experiment.yaml`** to set your experiment parameters: + + ```yaml + # @package _global_ + experiment: + exp_name: "my_custom_run" + seed: 42 + ``` + +## Training Execution + +To start a training run with the default settings defined in `configs/main_config.yaml`: + +```bash +uv run python scripts/train.py +``` + +### Using a Custom Experiment Configuration + +To run with your custom experiment file: + +```bash +uv run python scripts/train.py experiment=my_experiment +``` + +```bash +uv run python scripts/train.py ppo.learning_rate=0.001 ppo.num_envs=32 logging.track=true +``` + +## Evaluation During Training + +By default, the trainer saves checkpoints but does not evaluate them. To enable automatic headless evaluation of every saved checkpoint, set `evaluation.evaluate_checkpoints=true`: + +```bash +uv run python scripts/train.py evaluation.evaluate_checkpoints=true +``` + +## Reproducing Experiments + +For detailed steps on how to reproduce training runs, locate run configuration metadata, or reproduce our experiments using Weights & Biases (WandB), see the **[Results & Reproduction Guide](./reproduction.md)**. + +--- + +For more details on evaluation metrics and comparison tools, see [Checkpoint & Model Evaluation](./evaluation.md). + +For more details on tracking your experiments, see [Tracking & Monitoring](./tracking.md). diff --git a/docs/design/actor-critic.md b/docs/design/actor-critic.md new file mode 100644 index 0000000..8d03912 --- /dev/null +++ b/docs/design/actor-critic.md @@ -0,0 +1,121 @@ +# Actor-Critic Architecture + +To process observations into actions, our controllers utilize an Actor-Critic architecture. Because we use Proximal +Policy Optimization (PPO), the pipeline fundamentally requires separate networks for the policy (Actor) and the value +estimation (Critic). + +**Centralized Architecture (Baseline)** + +This pipeline treats the agent as a single entity and uses standard Proximal Policy Optimization (PPO). + +- Centralized Actor: Composed of two chained MLPs (Sensor $\rightarrow$ Motor) passing a hidden state between them. The + centralized sensor receives the concatenated global state vector of all limbs at once and processes it into a hidden + state. The centralized motor receives this hidden state and outputs the joint offsets for all actuators + simultaneously. This is mathematically equivalent to using one large MLP with hidden layers, but splitting makes the + implementation easier by allowing us to reuse the same components for the decentralized modules. +- Centralized Critic: Composed of two sequential MLPs (Feature Extractor $\rightarrow$ Critic). Because PPO evaluates + the state-value function, this network only receives the concatenated global state vector (no actions). It outputs a + single scalar estimating the expected future reward for the entire agent. + +Our policy and value networks use separate input networks/feature extractors as advised by the SEL3 course assistants and the blog. For continuous actions this should allow better learning at a small cost. + +```mermaid +graph TD + Obs([Global Observation]) + + Sens[Sensor] + Act[Motor] + OutAct([Action Distribution
mean, log_std]) + + Feat[Feature extractor] + Crit[Critic] + OutCrit([Value Estimate
scalar]) + + Obs --> Sens + Obs --> Feat + + Sens -->|"Hidden state"| Act + Feat -->|"Hidden state"| Crit + + Act --> OutAct + Crit --> OutCrit +``` + +**Decentralized Architecture** + +This pipeline utilizes the "Centralized Training with Decentralized Execution" principle, specifically the NerveNet-MLP +variant. + +- Decentralized Actor, split into three distinct models: + - Sensor: A local model at each node. It receives its local state plus the goal vector directly, processing them into + an initial hidden state. + - Propagator: Nodes synchronously compute and exchange messages with connected neighbors for $N$ steps to update + their hidden states. See [communication.md](./communication.md) for details. + - Motor: A local model uses its final updated hidden state to output the joint offset strictly for its own actuator. +- Centralized Critic: Composed of two sequential MLPs (Feature Extractor $\rightarrow$ Critic). During training, it + acts globally by taking the concatenated state vectors from all sensors to output a single, global state-value scalar + evaluating the entire agent's pose. + +To keep the implementation simple, we should use one critic per node in our architecture, but only a single, global +critic for all nodes at once, for the following reasons: + +1. Credit Assignment Problem (Ha, 2017): The MuJoCo simulator provides an overall reward based on the brittle star + movement progression, e.g. total distance travelled. Using an isolated critic for each node in the network would not + allow to determine which local action contributed to the global success. A global critic solves this by evaluating + the combined state of the agent at once. +2. Implementation simplicity: Building a second decentralized message-passing graph for the critic (NerveNet-2) would + require more coding. Using a standard MLP that concatenates all raw input vectors is much easier to program while + mathematically equivalent. + +```mermaid +graph TD + Obs([Local Observation]) + + Sens[Sensor] + Prop[Propagator] + Feat[Feature extractor] + + Mot[Motor] + Crit[Critic] + + OutMot([Action Distribution
mean, log_std]) + OutCrit([Value Estimate
scalar]) + + Obs --> Sens + Sens -->|"Hidden state"| Prop + Obs --> Feat + + Prop -->|"Hidden state"| Mot + + + Feat -->|"Hidden state"| Crit + + Mot --> OutMot + Crit --> OutCrit + + Prop -.->|"message passing"|Prop +``` + +## Implementation Details (Network Depth) + +Inspired by: [PPO Implementation Details](https://iclr-blog-track.github.io/2022/03/25/ppo-implementation-details/) + +The MLPs used in both pipelines are defined with specific hidden layer configurations to balance learning capability +and computational cost. As of right now, though this might change as we make progress in our experiments, we use: + +- Input Networks (Sensors & Feature Extractors): These networks map the raw state inputs to internal hidden states. + They are configured as standard dense networks with 3 hidden layers of 300 nodes each (`[300, 300, 300]`) and utilize `tanh` + activation functions. +- Output Networks (Motors, Actors & Critics): The final output models are intentionally kept shallow. The Actor + directly projects the hidden state to a continuous action distribution (`mean` and `log_std`) using a single dense + output layer (zero hidden layers) initialized orthogonally. The Critic functions similarly, mapping the hidden + representation to a single scalar value. + +Note: For the continuous action distributions outputted by the Motor, we explicitly use `mean` and `log_std` as advised +by previous research to maintain learning stability. + +**References** + +- Ha, D. (2017, October 29). A Visual Guide to Evolution Strategies. 大トロ ・ Machine Learning. [https://blog.otoro.net/2017/10/29/visual-evolution-strategies/](https://blog.otoro.net/2017/10/29/visual-evolution-strategies/) +- Schulman, John, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. ‘Proximal Policy Optimization Algorithms’. arXiv:1707.06347. Preprint, arXiv, 28 August 2017. [https://doi.org/10.48550/arXiv.1707.06347](https://doi.org/10.48550/arXiv.1707.06347). +- Wang, Tingwu, Renjie Liao, Jimmy Ba, and S. Fidler. ‘NerveNet: Learning Structured Policy with Graph Neural Networks’. Conference paper presented at International Conference on Learning Representations. 15 February 2018. [https://www.semanticscholar.org/paper/NerveNet:-Learning-Structured-Policy-with-Graph-Wang-Liao/249408527106d7595d45dd761dd53c83e5a02613](https://www.semanticscholar.org/paper/NerveNet:-Learning-Structured-Policy-with-Graph-Wang-Liao/249408527106d7595d45dd761dd53c83e5a02613). diff --git a/docs/design/communication.md b/docs/design/communication.md new file mode 100644 index 0000000..377de89 --- /dev/null +++ b/docs/design/communication.md @@ -0,0 +1,44 @@ +# Communication scheme (Message Passing) + +Remember our research question: + +> "What is the impact of different levels of controller modularity on learning speed, coordination, and fault tolerance +> (e.g. amputations) in brittle-star-like robots trained with Reinforcement Learning?" + +To test decentralized modularity (such as arm-level or segment-level controllers), the various modules *must* be able to +communicate with each other to achieve coordinated locomotion. This is accomplished through *message passing* in a Graph +Neural Network (GNN)-like architecture. Two prominent communication styles from the literature are N-step NerveNet (Wang +et al., 2018) and bottom-up top-down Shared Modular Policies (Huang et al., 2020). + +We have chosen to apply **one uniform communication style** across all modular architectures, specifically opting for +**N-step NerveNet**. + +## Rationale + +Initially, our idea was to equip arm-level controllers with NerveNet message passing and segment-level controllers with +SMP. However, we evaluated that this introduces a threat to the validity of our research question. If we observe +differences in performance, it would be impossible to determine whether the variance is caused by the *level of +modularity*, or by the difference in the message passing scheme. To purely compare modularity, the communication scheme +style must remain constant. + +Second, we decided that NerveNet is a better fit for our research. The morphology of our brittle star contains cycles at +the decentralized level (e.g., a ring of segments or arms around the body). NerveNet has proven to be robust for +arbitrary structures, including graphs with cycles. SMP inherently expects a tree structure for its bottom-up and +top-down pass. Applying SMP to a ring structure requires a workaround to break that cycle. + +## Limitations and alternatives + +Choosing NerveNet introduces a scalability issue as the morphology grows. In NerveNet, a message advances only one +segment or node per propagation step. When dealing with long arms (e.g., > 5 segments), this requires a large number of +propagation steps to transmit information from one tip of an arm to another. + +If we were to use SMP instead - which is possible - the inner states of nodes are shared across the entire graph in just +two passes. For very large or long morphologies, this would be much more scalable. + +By rejecting SMP, we accept that our model might learn slower or require more computational power for highly segmented, +extended morphologies. + +**References** + +- Wang, Tingwu, Renjie Liao, Jimmy Ba, and S. Fidler. ‘NerveNet: Learning Structured Policy with Graph Neural Networks’. Conference paper presented at International Conference on Learning Representations. 15 February 2018. [https://www.semanticscholar.org/paper/NerveNet:-Learning-Structured-Policy-with-Graph-Wang-Liao/249408527106d7595d45dd761dd53c83e5a02613](https://www.semanticscholar.org/paper/NerveNet:-Learning-Structured-Policy-with-Graph-Wang-Liao/249408527106d7595d45dd761dd53c83e5a02613). +- Huang, Wenlong, Igor Mordatch, and Deepak Pathak. ‘One Policy to Control Them All: Shared Modular Policies for Agent-Agnostic Control’. arXiv:2007.04976. Preprint, arXiv, 9 July 2020. [https://doi.org/10.48550/arXiv.2007.04976](https://doi.org/10.48550/arXiv.2007.04976). diff --git a/docs/design/controllers.md b/docs/design/controllers.md new file mode 100644 index 0000000..2dd59d9 --- /dev/null +++ b/docs/design/controllers.md @@ -0,0 +1,28 @@ +# Levels of modularity and topology + +The brittle star can be controlled at different levels. A monolithic controller processes all inputs and outputs at +once, whereas modular controllers divide the brains across the body, inspired by the biology of brittle stars. + +We define four architectures to compare: + +1. **Centralized, monolithic**: A single Multi Layer Perceptron per robot that receives all observations and outputs + all actions. +2. **Fully connected arm-level**: Each arm contains an MLP that processes the inputs for that arm, an MLP that processes + the communicated inner-states, and an MLP that outputs the actions for that arm. One policy for these MLPs is shared + across the arms. The controllers in each arm are connected to each other and form a fully connected graph. There is + no central disk, but the controllers are fully connected. +3. **Ring arm-level**: Identical setup to the fully connected arm-level, but the controllers are connected in a ring + structure. This setup is considered less centralized than the fully connected graph. +4. **Segment-level**: Each segment contains the three MLPs discussed above. The base segments, attached to the body, + form a ring structure, with the remaining segments attached as extended "strings". Segments can only communicate with + segments that are physically connected to it. + +## Rationale + +To fairly compare decentralized modularity against centralized control, the decentralized models should not be allowed +to contain a central organ acting as a bottleneck or coordinator. By removing the central disk in the decentralized +models and replacing it with a ring topology, we closely approximate the biological reality of the brittle star and test +a decentralized morphology. + +The fully connected graph functions as an intermediate step in between a fully centralized and a decentralized ring. We +use it to test whether our models scale to more complex structures. diff --git a/docs/design/input_action_spaces.md b/docs/design/input_action_spaces.md new file mode 100644 index 0000000..ddb7a55 --- /dev/null +++ b/docs/design/input_action_spaces.md @@ -0,0 +1,164 @@ +# Input (state) and output (action) spaces + +To effectively learn locomotion and navigation, the agent requires a well-defined observation space (inputs) and action +space (outputs). The control models map these observations directly to physical movements. + +**Inputs (state space)** + +The observation space provides the agent with its current physical state and its navigational objective. With a +decentralized control architecture in mind, we divide these inputs into global and local states. + +Global inputs, always broadcasted to all nodes: + +- Vertical orientation/tilt: A single, simplified metric representing the tilt/vertical alignment of the agent's + central body/disk, a.k.a. the deviation from the global Z-axis. Its value is derived from the environment's raw disk + rotation 3D vector $[roll, pitch, yaw]$: + $$ + tilt = sqrt(roll^2 + pitch^2) + $$ +- Goal vector: A 2D unit vector representing the *egocentric* direction to the target. A value of $[1.0, 0.0]$ + indicates that the target is directly in front of the agent (angle 0). + +Local inputs, routed directly to specific nodes: + +- Joint positions: The current angles of all joints within the morphology. +- Joint velocities: The current angular velocities of the joints. +- Joint actuator forces: The physical forces currently exerted at each specific joint. +- Segment contact: These values indicate whether each physical segment of the agent is currently touching the ground. + +**Outputs (action space)** + +The action space defines how the agent interacts with the environment. + +- Joint offsets: *absolute* target positions (offsets) for the joints, i.e. the exact angle the joint should move to. + +## Normalization and Scaling + +Both the input (observation) and output (action) spaces are rescaled to the range **$[-1, 1]$**. + +For the input space, all raw physical values (angles, velocities, forces, distances) are normalized based on their +defined physical bounds. If a value exceeds these bounds during simulation, it is clipped to the $[-1, 1]$ range. + +For the output space, the neural network's tanh-activated outputs (which naturally fall in $[-1, 1]$) are linearly +mapped to the physical joint limits defined in the robot's morphology. + +## Rationale + +When designing the state space, we must ask: *Could a human operator perform this task given only these inputs?* + +- Inclusion of Joint Velocities: Because our control models do not inherently possess memory of previous timesteps, + providing only the joint position is insufficient to determine the direction a limb is currently moving. By + explicitly including joint velocities, the agent can immediately infer momentum and movement direction without + needing to memorize past states. +- Absolute Joint Offsets: The physical Brittle Star robot relies on servo motors (if we were to build this simulated + robot), which are inherently position-controlled devices. (Continuous rotation servos exist, but they are less + commonly used for joints.) If our network outputted continuous torques (forces), a significant portion of the + reinforcement learning process would be wasted on learning low-level PID control dynamics (i.e., how much force to + apply to hold a position). Abstracting this away forces the learning algorithm to focus entirely on higher-level gait + generation and locomotion. +- Simplified vertical orientation: We drop the full 3D spatial rotation and angular velocity arrays in favor of a + single vertical orientation metric (tilt). For a brittle star moving accross a flat plane, this metric is sufficient + for the agent to sense if it is losing balance or flipping over. +- Force representation: We strictly retain the joint actuator forces and drop the more generic actuator force. Forces + that are explicitly tied to individual joints are significantly easier to route into decentralized, local limb nodes, + which is necessary for our message-passing architecture. +- Goal Vector (Distance + Angle): Providing only the scalar "distance to the goal" as an input is akin to blindfolding + the robot and asking it to find a target by playing "hot or cold." By providing a full vector, the agent knows + exactly where the target is relative to its current orientation, allowing for directed and efficient locomotion. + + The goal representation is explicitly divided into a directional vector and a scalar distance. Keeping the direction + as a normalized unit vector bounds the values to the $[-1, 1]$ range, which stabilizes neural network training. + Providing only a scalar "distance to the goal" would force the agent to learning localized searching behaviors (e.g. + random walks or spiraling) to deduce the direction, drastically increasing the difficulty of the learning task. + + **NOTE:** We later dropped the "distance to vector", switching to only a direction as the input. Our reasoning is + the agent should always move towards the goal (it should not learn to stop at the goal), which allows for this + simplification that decreases the model input size. + + The environment provides a raw `unit_xy_direction_to_target` (global), which we transform into a calculated + `robot_direction_to_target` (egocentric) before passing it to the MLPs. This vector consists of the X and Y + direction, where a value of $[1.0, 0.0]$ (mapping to an angle of $0$) means the robot is facing directly towards the + target. +- Contact sensors: Segment contact detects external ground interaction and is biologically vital for timing gait + transitions. +- Zero-Centered Rescaling ($[-1, 1]$): Using a zero-centered range is standard best practice for continuous control + tasks. It provides several mathematical and physical advantages: + - Improved Gradient Flow: Neural networks optimize faster when inputs are zero-centered. If all inputs were positive + (e.g., $[0, 1]$), the gradients during backpropagation would be forced to the same sign, causing inefficient + "zig-zag" weight updates. + - Meaningful Neutral State: In robotics, $0.0$ naturally represents a resting state (zero velocity, centered + position, no force). In a $[-1, 1]$ system, this physical rest maps to a neutral $0.0$ signal in the network. + This also correctly communicaties a "neutral/dead" signal for amputated limbs that are padded with $0.0$ values. + +Specifically, we do not include some available inputs: + +- Global position: Absolute spatial coordinates can cause the agent to overfit to a specific coordinate frame or map, + rather than learning general, adaptable locomotion strategies. + +## Limitations and alternatives + +Alternative state and action formulations include: + +- Torque-based continuous control: In many continuous control tasks (like standard MuJoCo benchmarks), actions + represent continuous torques applied to joints. While this provides more granular, low-level physical control, it + heavily complicates training and does not align well with the physical reality of servo-driven hardware. +- Recurrent Neural Networks (RNNs) / Frame Stacking: Instead of explicitly passing velocities in the state space, the + network could infer momentum by observing a history of past states. Using RNNs or frame stacking allows the agent to + build an internal memory of movement. However, this significantly increases architectural complexity and training + time compared to explicitly providing the velocity data. +- Scalar Goal Distance: Giving the agent only the scalar distance to the target would force it to learn a localized + searching behavior (e.g., spiraling or random walks) to determine the correct direction. While biologically plausible + for simpler organisms following chemical gradients, it drastically increases the difficulty of the learning task. +- **$[0, 1]$ Rescaling:** While some domains (like computer vision) use $[0, 1]$ scaling, it is generally avoided in + robotics. Scaling to $[0, 1]$ would mean that a resting joint (velocity = 0) maps to an input of $0.5$. This + constant positive bias forces the network to waste capacity learning to ignore or subtract this baseline signal just + to stand still. Furthermore, it breaks the "dead signal" interpretation of zero-padding used for amputations. + +## MuJoCo + +This is what the filtered input vectors look like in MuJoCo, with $J$ joints and $S$ segments: + +- `joint_position`: shape=(J,), dtype=float64 +- `joint_velocity`: shape=(J,), dtype=float64 +- `joint_actuator_force`: shape=(J,), dtype=float64 +- `segment_contact`: shape=(S,), dtype=float64 +- `robot_direction_to_target`: shape=(2,), dtype=float64, egocentric +- `disk_z_tilt`: shape=(1,), dtype=float64, derived from `disk_rotation` + +This brings the entire input space down to $3J + S + 4$ float64's, compared to $4J + S + 15$ float64's for the +unfiltered inputs. + +For reference, these are all the inputs that are available in the MuJoCo environment: + +``` +obs keys: ['joint_position', 'joint_velocity', 'joint_actuator_force', 'actuator_force', 'disk_position', 'disk_rotation', 'disk_linear_velocity', 'disk_angular_velocity', 'tendon_position', 'tendon_velocity', 'segment_contact', 'unit_xy_direction_to_target', 'xy_distance_to_target'] + +raw observations dict: +{'joint_position': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]), + 'joint_velocity': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]), + 'joint_actuator_force': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]), + 'actuator_force': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]), + 'disk_position': array([0. , 0. , 0.11]), + 'disk_rotation': (0.0, -0.0, 0.0), + 'disk_linear_velocity': array([0., 0., 0.]), + 'disk_angular_velocity': array([0., 0., 0.]), + 'tendon_position': array([], dtype=float64), + 'tendon_velocity': array([], dtype=float64), + 'segment_contact': array([0., 0., 0., 0., 0., 0.]), + 'unit_xy_direction_to_target': array([-0.95333378, -0.30191837]), + 'xy_distance_to_target': array([3.])} + +(shapes) +joint_position: shape=(12,), dtype=float64, size=12 +joint_velocity: shape=(12,), dtype=float64, size=12 +joint_actuator_force: shape=(12,), dtype=float64, size=12 +actuator_force: shape=(12,), dtype=float64, size=12 +disk_position: shape=(3,), dtype=float64, size=3 +disk_rotation: shape=(3,), dtype=float64, size=3 +disk_linear_velocity: shape=(3,), dtype=float64, size=3 +disk_angular_velocity: shape=(3,), dtype=float64, size=3 +tendon_position: shape=(0,), dtype=float64, size=0 +tendon_velocity: shape=(0,), dtype=float64, size=0 +segment_contact: shape=(6,), dtype=float64, size=6 +xy_distance_to_target: shape=(1,), dtype=float64, size=1 +``` diff --git a/docs/design/learning_algorithm.md b/docs/design/learning_algorithm.md new file mode 100644 index 0000000..afee110 --- /dev/null +++ b/docs/design/learning_algorithm.md @@ -0,0 +1,30 @@ +# Reinforcement Learning Algorithm + +To control the *continuous* action space (the joints of the robot) based on sensor data, we require a reliable +Reinforcement Learning (RL) algorithm or optimization strategy. + +We have chosen **Proximal Policy Optimization (PPO)** (Schulman et al., 2017). + +## Rationale + +PPO is an on-policy algorithm known for its stability and robustness (safe training without excessive variance). More +importantly, it requires relatively little hyperparameter tuning compared to other algorithms. Since NerveNet was +successfully trained using PPO (Wang et al., 2018), selecting PPO significantly reduces the risk of convergence +failures. + +## Limitations and alternatives + +Alternative learning algorithms include: + +- **Twin Delayed DDPG (Fujimoto et al., 2018)**: TD3 is a strong off-policy alternative used in the SMP paper (Huang et + al., 2020). It is highly sample-efficient and reportedly excels at zero-shot adaptations. However, this approach would + be more complex and error-prone than with PPO. +- **Evolution strategies (ES)**: Evolution strategies are useful for optimizing Central Pattern Generators (CPGs), e.g. + CMA-ES, OpenAI-ES. While this method is easier to distribute and parallelize, ES typically scales worse with + exceptionally large observation spaces compared to gradient-based RL methods like PPO. + +**References** + +- Fujimoto, Scott, Herke Hoof, and David Meger. ‘Addressing Function Approximation Error in Actor-Critic Methods’. Proceedings of the 35th International Conference on Machine Learning, 3 July 2018, 1587-96. [https://proceedings.mlr.press/v80/fujimoto18a.html](https://proceedings.mlr.press/v80/fujimoto18a.html). +- Huang, Wenlong, Igor Mordatch, and Deepak Pathak. ‘One Policy to Control Them All: Shared Modular Policies for Agent-Agnostic Control’. arXiv:2007.04976. Preprint, arXiv, 9 July 2020. [https://doi.org/10.48550/arXiv.2007.04976](https://doi.org/10.48550/arXiv.2007.04976). +- Schulman, John, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. ‘Proximal Policy Optimization Algorithms’. arXiv:1707.06347. Preprint, arXiv, 28 August 2017. [https://doi.org/10.48550/arXiv.1707.06347](https://doi.org/10.48550/arXiv.1707.06347). diff --git a/docs/design/reward_function.md b/docs/design/reward_function.md new file mode 100644 index 0000000..afbf18b --- /dev/null +++ b/docs/design/reward_function.md @@ -0,0 +1,32 @@ +# Reward function and observation space + +The robot needs to know whether its movements contribute to the ultimate goal of locomotion towards a target. Sensor +inputs must be distributed fairly to guarantee an objective comparison between different architectures. + +- The distance from the robot to the target and/or the light intensity are treated as global inputs. +- Positions and joints, which are normalized to floating-point values between 0 and 1, are considered local inputs. +- The reward function is centered around minimizing the distance to the goal or maximizing the movement towards the goal + within a finite number of timesteps $T$. +- To motivate efficient movement, the amount of timesteps taken to reach the goal will be used as penalty. +- An extra penalty based on movement relative to the current step and the previous is used to penalize a movement away from the target. + +## From reward to PPO + +The resulting reward is passed to our PPO library. Our critic network (value function) predicts how good our eventual reward will be for the current state, this value is combined with the reward from the reward function to get advantages. These advantages are then used to calculate the losses to update both our critic and actor pipeline. + +## Rationale + +Using a light source (or a gradient) is biologically plausible for many simple organisms. By normalizing all signals +between 0 and 1, PPO training is highly stabilized. The timesteps must be finite to reset the environment in a timely +manner if the policy gets stuck in a local minimum. + +## Limitations and alternatives + +Providing global information to all individual decentralized segments can be considered biologically cheating or +practically infeasible once the robot would be physically built. Some sensory input cannot be put in each joint, for +example. + +The alternative is to provide the global input to the outermost segments of the arms, or a specific set of segments +assigned with this functionality. The network would then have to learn to propagate this signal throughout the body via +message passing. While biologically more accurate, this drastically complicates the learning process. We have written +this down as potential future research. diff --git a/docs/javascripts/katex.js b/docs/javascripts/katex.js new file mode 100644 index 0000000..956b86a --- /dev/null +++ b/docs/javascripts/katex.js @@ -0,0 +1,16 @@ +const renderMath = (el) => { + renderMathInElement(el, { + delimiters: [ + { left: "$$", right: "$$", display: true }, + { left: "$", right: "$", display: false }, + { left: "\\(", right: "\\)", display: false }, + { left: "\\[", right: "\\]", display: true } + ], + }); +}; + +if (typeof document$ !== "undefined") { + document$.subscribe(({ body }) => renderMath(body)); +} else { + document.addEventListener("DOMContentLoaded", () => renderMath(document.body)); +} diff --git a/env/hpc/modules.txt b/env/hpc/modules.txt new file mode 100644 index 0000000..2ea2c6e --- /dev/null +++ b/env/hpc/modules.txt @@ -0,0 +1,4 @@ +GCCcore/13.3.0 +Python/3.12.3-GCCcore-13.3.0 +FFmpeg/7.0.2-GCCcore-13.3.0 +Hydra/1.3.2-GCCcore-13.3.0 diff --git a/env/hpc/requirements.txt b/env/hpc/requirements.txt new file mode 100644 index 0000000..c85ce2d --- /dev/null +++ b/env/hpc/requirements.txt @@ -0,0 +1,20 @@ +biorobot==0.4.2 +cleanrl>=0.4.8 +evosax==0.2.0 +flax>=0.12.2 +gymnasium>=1.2.3 +ipykernel==7.2.0 +jax[cuda13]==0.9.0.1 +numpy>=2.0.0 +protobuf>=5.0.0 +warp-lang +mujoco-warp +matplotlib==3.10.8 +mediapy==1.2.6 +optax>=0.2.6 +pyopengl>=3.1.10 +pyopengl-accelerate>=3.1.10 +pyyaml>=6.0 +hydra-core>=1.3.2 +wandb==0.24.2 +torch>=2.4.0 diff --git a/flake.nix b/flake.nix index c766ec9..bdcdf41 100644 --- a/flake.nix +++ b/flake.nix @@ -29,6 +29,10 @@ # Editor of your choice (nix-jetbrains-plugins.lib.buildIdeWithPlugins pkgs "pycharm" pluginList) ]; + + shellHook = '' + uv run pre-commit install + ''; }; }); } diff --git a/mkdocs.yml b/mkdocs.yml new file mode 100644 index 0000000..cc60194 --- /dev/null +++ b/mkdocs.yml @@ -0,0 +1,41 @@ +site_name: Brittle Star Project +theme: + name: material + +nav: + - Home: README.md + - Design & Architecture: + - Actor-Critic Architecture: design/actor-critic.md + - Communication Scheme: design/communication.md + - Modularity & Topology: design/controllers.md + - Input & Action Spaces: design/input_action_spaces.md + - Reinforcement Learning Algorithm: design/learning_algorithm.md + - Reward Function & Observation Space: design/reward_function.md + - API Reference: + - Brittle Star Environment: api/environment.md + - Training Models: api/training.md + - Tracking & Monitoring: api/tracking.md + - Checkpoint & Model Evaluation: api/evaluation.md + - Interactive Simulation & Visualization: api/simulation.md + - Analysis & Plotting Tools: api/analysis.md + - Results & Reproduction: api/reproduction.md + - HPC Guide: HPC.md + - Contribution Guidelines: CONTRIBUTING.md + - Development Guide: DEVELOPMENT.md + +markdown_extensions: + - pymdownx.superfences: + custom_fences: + - name: mermaid + class: mermaid + format: !!python/name:pymdownx.superfences.fence_code_format + - pymdownx.arithmatex: + generic: true + +extra_css: + - https://unpkg.com/katex@0/dist/katex.min.css + +extra_javascript: + - javascripts/katex.js + - https://unpkg.com/katex@0/dist/katex.min.js + - https://unpkg.com/katex@0/dist/contrib/auto-render.min.js diff --git a/pyproject.toml b/pyproject.toml index a44dd05..df3c6a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,22 +1,84 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + [project] name = "2026sel3-project" version = "0.1.0" description = "Add your description here" readme = "README.md" -requires-python = ">=3.12, <3.13" +requires-python = ">= 3.12, < 3.13" dependencies = [ "biorobot==0.4.2", + "cleanrl>=0.4.8", "evosax==0.2.0", + "flax>=0.12.2", + "gymnasium>=1.2.3", "ipykernel==7.2.0", - "jax[cuda13]==0.9.0.1", + "jax==0.9.0.1", + "numpy>=2.0.0", + "protobuf>=5.0.0", + "warp-lang", + "mujoco-warp", "matplotlib==3.10.8", "mediapy==1.2.6", + "optax>=0.2.6", "pyopengl>=3.1.10", "pyopengl-accelerate>=3.1.10", + "pyyaml>=6.0", + "hydra-core>=1.3.2", "wandb==0.24.2", + "torch>=2.4.0", +] + +[project.optional-dependencies] +cuda = [ + "jax[cuda13]==0.9.0.1", +] +analysis = [ + "tensorboard", +] +evaluation = [ + "imageio>=2.35.0", + "imageio-ffmpeg>=0.5.1", ] [dependency-groups] dev = [ + "pre-commit>=4.0.0", + "pytest>=8.0.0", "ruff>=0.15.2", ] + +[tool.hatch.build.targets.wheel] +packages = ["src/brittle_star_project", "src/experiment_logger"] + +[tool.mypy] +mypy_path = "src" +check_untyped_defs = false +warn_return_any = false + +[[tool.mypy.overrides]] +module = [ + "jax.*", + "flax.*", + "wandb.*", + "torch.*", + "mujoco.*", + "mujoco_warp.*", + "optax.*", + "tyro.*", + "biorobot.*", + "gymnasium.*", + "matplotlib.*", + "mediapy.*", + "matplotlib.*", + "mediapy.*", + "pytest.*", + "tensorboard.*", + "tqdm.*", + "numpy.*", + "yaml.*", + "moojoco.*" +] +ignore_missing_imports = true diff --git a/ruff.toml b/ruff.toml index d3c196a..ad24132 100644 --- a/ruff.toml +++ b/ruff.toml @@ -1,3 +1,7 @@ +line-length = 100 +exclude = ["wandb"] + +[lint] extend-select = [ # "PLC0103", # invalid-name # "PLC0104", # disallowed-name @@ -40,10 +44,10 @@ extend-select = [ "PLC2401", # non-ascii-name # "PLC2403", # non-ascii-module-import # "PLC2503", # bad-file-encoding - "PLC2801", # unnecessary-dunder-call + # "PLC2801", # unnecessary-dunder-call (requires preview) # "PLC3001", # unnecessary-lambda-assignment "PLC3002", # unnecessary-direct-lambda-call - "E999", # syntax-error + # "E999", # syntax-error (removed from ruff) # "PLE0011", # unrecognized-inline-option # "PLE0013", # bad-plugin-value # "PLE0014", # bad-configuration-section @@ -130,7 +134,7 @@ extend-select = [ # "PLE1137", # unsupported-assignment-operation # "PLE1138", # unsupported-delete-operation # "PLE1139", # invalid-metaclass - "PLE1141", # dict-iter-missing-items + # "PLE1141", # dict-iter-missing-items (requires preview) "PLE1142", # await-outside-async # "PLE1143", # unhashable-member # "PLE1144", # invalid-slice-step @@ -168,7 +172,7 @@ extend-select = [ # "PLE3102", # positional-only-arguments-expected # "PLE3701", # invalid-field-call # "PLE4702", # modified-iterating-dict - "PLE4703", # modified-iterating-set + # "PLE4703", # modified-iterating-set (requires preview) # "PLF0001", # fatal # "PLF0002", # astroid-error # "PLF0010", # parse-error @@ -210,7 +214,7 @@ extend-select = [ # "PLW0238", # unused-private-member # "PLW0239", # overridden-final-method # "PLW0240", # subclassed-final-class - "PLW0244", # redefined-slots-in-subclass + # "PLW0244", # redefined-slots-in-subclass (requires preview) "PLW0245", # super-without-brackets # "PLW0246", # useless-parent-delegation # "PLW0301", # unnecessary-semicolon @@ -274,7 +278,7 @@ extend-select = [ "PLW1508", # invalid-envvar-default "PLW1509", # subprocess-popen-preexec-fn # "PLW1510", # subprocess-run-check - "PLW1514", # unspecified-encoding + # "PLW1514", # unspecified-encoding (requires preview) # "PLW1515", # forgotten-debug-statement # "PLW1518", # method-cache-max-size-none "PLW2101", # useless-with-lock @@ -298,7 +302,7 @@ extend-select = [ # "PLW4906", # deprecated-attribute ] -ignore = [ +extend-ignore = [ # "PLC0116", # missing-function-docstring # "PLC0200", # consider-using-enumerate # "PLC0305", # trailing-newlines @@ -346,7 +350,7 @@ ignore = [ # "PLR1705", # no-else-return # "PLR1706", # consider-using-ternary # "PLR1707", # trailing-comma-tuple - "PLR1708", # stop-iteration-return + # "PLR1708", # stop-iteration-return (deprecated) # "PLR1709", # simplify-boolean-expression # "PLR1710", # inconsistent-return-statements "PLR1711", # useless-return @@ -391,4 +395,5 @@ ignore = [ # "PLW1404", # implicit-str-concat ] - +[lint.per-file-ignores] +"__init__.py" = ["F401"] diff --git a/scripts/analysis/README.md b/scripts/analysis/README.md new file mode 100644 index 0000000..a99f952 --- /dev/null +++ b/scripts/analysis/README.md @@ -0,0 +1,27 @@ +# Experiment Analysis Tools + +This directory contains scripts for post-processing and analyzing experiment results, including TensorBoard logs and saved model weights. + +## Scripts + +### 1. `explore_tensorboard.py` +A CLI tool to summarize TensorBoard `tfevents` files without a GUI. + +**Key Features:** +- Displays last values, min, max, and step counts for all scalar metrics. +- Calculates total run duration and estimated completion percentage. +- Exports granular scalar data to CSV for analysis in Excel/Pandas. + +**Usage:** +```bash +# General usage +python explore_tensorboard.py + +# Exporting data +python explore_tensorboard.py --csv data.csv +``` + +**Requirements:** +- `pandas` +- `tensorboard` +- `tensorflow-cpu` (or `tensorflow`) diff --git a/scripts/analysis/explore_tensorboard.py b/scripts/analysis/explore_tensorboard.py new file mode 100644 index 0000000..0ed1027 --- /dev/null +++ b/scripts/analysis/explore_tensorboard.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python3 +""" +Reproducible CLI tool to explore TensorBoard logs. +Designed for both local development and HPC diagnostics. + +Requirements: + pip install tensorboard + +Usage: + python explore_tensorboard.py [--csv output.csv] +""" + +import argparse +import os +import sys +import csv + +try: + from tensorboard.backend.event_processing import event_accumulator +except ImportError: + print("Error: Missing dependency. Please run: pip install tensorboard") + sys.exit(1) + + +def explore_run(log_dir): + """ + Extracts and displays a summary of scalar metrics from a TensorBoard log directory. + """ + print(f"\n{'=' * 20} Exploring Run {'=' * 20}") + print(f"Directory: {log_dir}") + print(f"{'=' * 55}\n") + + if not os.path.exists(log_dir): + print(f"Error: Directory '{log_dir}' does not exist.") + return None + + # Initialize EventAccumulator + # size_guidance=0 loads all data points for each tag. + ea = event_accumulator.EventAccumulator( + log_dir, + size_guidance={ + event_accumulator.SCALARS: 0, + event_accumulator.TENSORS: 0, + }, + ) + + print("Loading event files (this may take a moment for large runs)...") + ea.Reload() + + tags = ea.Tags() + scalar_tags = tags.get("scalars", []) + + if not scalar_tags: + print("No scalar metrics found in this directory.") + return None + + print(f"Found {len(scalar_tags)} scalar metrics.\n") + + data = {} + summary = [] + + # Process scalar values + for tag in scalar_tags: + events = ea.Scalars(tag) + if not events: + continue + + values = [e.value for e in events] + last_event = events[-1] + data[tag] = values + + summary.append( + { + "Metric": tag, + "Steps": len(events), + "Last Value": f"{last_event.value:.4f}", + "Max": f"{max(values):.4f}", + "Min": f"{min(values):.4f}", + } + ) + + # Display summary table formatted manually + summary = sorted(summary, key=lambda x: x["Metric"]) + print(f"{'Metric':<30} {'Steps':>10} {'Last':>12} {'Max':>12} {'Min':>12}") + print("-" * 80) + for row in summary: + print( + f"{row['Metric']:<30} {row['Steps']:>10} {row['Last Value']:>12} " + f"{row['Max']:>12} {row['Min']:>12}" + ) + + # Calculate and display global metadata + if "charts/SPS" in data: + sps_events = ea.Scalars("charts/SPS") + if len(sps_events) > 1: + total_duration_hours = (sps_events[-1].wall_time - sps_events[0].wall_time) / 3600 + print(f"\nTotal Recorded Duration: {total_duration_hours:.2f} hours") + + # Estimate completion if total_timesteps is available in hyperparameters + try: + hp_tags = [t for t in tags.get("tensors", []) if "hyperparameters" in t] + if hp_tags: + hp_event = ea.Tensors(hp_tags[0])[0] + hp_text = hp_event.tensor_proto.string_val[0].decode("utf-8") + if "total_timesteps" in hp_text: + for line in hp_text.split("\n"): + if "total_timesteps" in line: + target = int(line.split("|")[2].strip()) + current = ea.Scalars(scalar_tags[0])[-1].step + percent = (current / target) * 100 + print(f"Progress: {current:,} / {target:,} steps ({percent:.1f}%)") + except Exception: + pass + + return data + + +def main(): + parser = argparse.ArgumentParser(description="Reproducible TensorBoard exploration tool.") + parser.add_argument("log_dir", help="Path to the TensorBoard run directory.") + parser.add_argument("--csv", help="Optional: Path to export scalar data to CSV.", default=None) + + args = parser.parse_args() + + scalar_data = explore_run(args.log_dir) + + if args.csv and scalar_data: + # Reloading for wall_time and steps + ea = event_accumulator.EventAccumulator(args.log_dir).Reload() + with open(args.csv, mode="w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=["tag", "step", "value", "wall_time"]) + writer.writeheader() + for tag in scalar_data.keys(): + for e in ea.Scalars(tag): + writer.writerow( + {"tag": tag, "step": e.step, "value": e.value, "wall_time": e.wall_time} + ) + + print(f"\nData exported to: {args.csv}") + + +if __name__ == "__main__": + main() diff --git a/scripts/compare_models.py b/scripts/compare_models.py new file mode 100644 index 0000000..b3e8338 --- /dev/null +++ b/scripts/compare_models.py @@ -0,0 +1,182 @@ +"""Compare multiple trained policies across shared evaluation conditions. + +For each model listed in evaluation.comparison_models, this script runs +`comparison_num_episodes` headless rollouts (seeded sequentially from +`comparison_base_seed`) and writes a results CSV to `comparison_output_csv`. + +Results include two metrics per episode: +- `eval_return` — shaped reward (same function used during training) +- `max_velocity` — approximated as initial_xy_dist / steps taken + +Usage: + # With the default evaluation config + python scripts/compare_models.py evaluation=poster + + # Override the output path on the fly + python scripts/compare_models.py evaluation=poster \\ + evaluation.comparison_output_csv=metrics/quick_comparison.csv +""" + +from __future__ import annotations + +import csv +import logging +import time +from pathlib import Path + +import hydra +from omegaconf import DictConfig, OmegaConf + +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.configs.register_configs import register_configs +from brittle_star_project.evaluation import build_eval_env +from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs +from brittle_star_project.evaluation.rollout import rollout_headless + +_FIELDNAMES = [ + "model_path", + "architecture", + "arm_0", + "arm_1", + "arm_2", + "arm_3", + "arm_4", + "num_active_arms", + "seed", + "reached_target", + "episode_length", + "eval_return", + "initial_target_distance", + "final_xy_dist", + "approx_max_velocity", +] + + +def _approx_max_velocity(result) -> float | None: + """Approximate max velocity as distance covered per step. + + This is a rough upper bound: (initial_dist - final_dist) / steps. + """ + if result.initial_target_distance is None or result.final_xy_dist is None or result.length <= 0: + return None + dist_covered = result.initial_target_distance - result.final_xy_dist + return dist_covered / result.length + + +@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") +def main(dict_cfg: DictConfig) -> None: + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + logger = logging.getLogger(__name__) + + cfg: BrittleStarConfig = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg) + ) + eval_cfg = cfg.evaluation + + model_paths = [str(p) for p in eval_cfg.comparison_models] + if not model_paths: + raise ValueError( + "evaluation.comparison_models is empty. " + "Add at least one model path in your evaluation config." + ) + + base_seed = int(eval_cfg.comparison_base_seed) + num_episodes = int(eval_cfg.comparison_num_episodes) + max_steps = int(eval_cfg.eval_max_steps) + + seeds = list(range(base_seed, base_seed + num_episodes)) + + output_path = Path(hydra.utils.to_absolute_path(eval_cfg.comparison_output_csv)) + output_path.parent.mkdir(parents=True, exist_ok=True) + + logger.info( + f"Comparing {len(model_paths)} models over {num_episodes} episodes " + f"(seeds {seeds[0]}–{seeds[-1]})." + ) + logger.info(f"Results will be written to: {output_path}") + + with open(output_path, "w", newline="") as csv_file: + writer = csv.DictWriter(csv_file, fieldnames=_FIELDNAMES) + writer.writeheader() + + for model_path_str in model_paths: + model_path = Path(hydra.utils.to_absolute_path(model_path_str)) + logger.info(f"Evaluating model: {model_path.name}") + + try: + metadata = load_metadata(model_path) + except FileNotFoundError as e: + logger.warning(f"Skipping model — {e}") + continue + + training = metadata_to_configs(metadata) + + # Determine morphologies to evaluate + # If comparison_morphologies is empty, use the model's training morphology + morphologies = [None] + if eval_cfg.comparison_morphologies: + morphologies = [ + Path(hydra.utils.to_absolute_path(m)) for m in eval_cfg.comparison_morphologies + ] + + for morph_path in morphologies: + morph_label = morph_path.name if morph_path else "training" + logger.info(f" Morphology: {morph_label}") + + bundle = build_eval_env( + model_path=model_path, + training=training, + metadata=metadata, + morphology_override_path=morph_path, + ) + + for seed in seeds: + t0 = time.time() + result = rollout_headless( + env=bundle.env, + policy=bundle.policy, + seed=seed, + max_steps=max_steps, + action_low=bundle.action_low, + action_high=bundle.action_high, + action_mask=bundle.action_mask, + ) + elapsed = time.time() - t0 + + velocity = _approx_max_velocity(result) + + logger.debug( + f" seed={seed:3d} | " + f"reached={str(result.reached_target):<5} | " + f"return={result.return_:+8.3f} | " + f"steps={result.length:4d} | " + f"({elapsed:.1f}s)" + ) + + row = { + "model_path": model_path_str, + "architecture": bundle.architecture, + "num_active_arms": bundle.num_active_arms, + "seed": seed, + "reached_target": result.reached_target, + "episode_length": result.length, + "eval_return": result.return_, + "initial_target_distance": result.initial_target_distance, + "final_xy_dist": result.final_xy_dist, + "approx_max_velocity": velocity, + } + # Add per-arm segments + for i, segs in enumerate(bundle.segments_per_arm): + row[f"arm_{i}"] = segs + + writer.writerow(row) + csv_file.flush() + + bundle.env.close() + + logger.info(f"Done. Results saved to {output_path}") + + +if __name__ == "__main__": + register_configs() + main() diff --git a/scripts/evaluate_checkpoints.py b/scripts/evaluate_checkpoints.py new file mode 100644 index 0000000..a63796b --- /dev/null +++ b/scripts/evaluate_checkpoints.py @@ -0,0 +1,264 @@ +"""Re-evaluate saved checkpoints from a completed training run using MJX. + +This script scans the checkpoint directory of a training run (the `checkpoints/` +folder inside a Hydra output directory), loads each `.flax` checkpoint, runs +one deterministic evaluation episode with `build_eval_rollout_fn`, and appends +the result to the run's `metrics/checkpoint_evaluation.csv`. + +It is intended for post-training analysis when per-checkpoint evaluation was not +enabled during training (`evaluate_checkpoints: false`). + +Usage: + python scripts/evaluate_checkpoints.py \ + simulation.model_path=runs/2024-01-01/12-00-00/final_model.flax \ + evaluation.eval_max_steps=5000 \ + evaluation.eval_seed=0 + +The script resolves the run directory from `simulation.model_path`, discovers +all `*.flax` checkpoints under `checkpoints/`, and evaluates them in order. +""" + +from __future__ import annotations +from brittle_star_project.MLPs.mlps import ( + Actor, + GenericDenseLayersWithActivation, + MessagePasser, +) +from brittle_star_project.MLPs.adjancency_builder import build_adjacency +from brittle_star_project.environment import MorphMode +from brittle_star_project.MLPs.routing import apply_per_node +import logging +import re +from pathlib import Path + +import hydra +import jax +import numpy as np +import jax.numpy as jnp + +from omegaconf import DictConfig, OmegaConf + +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.configs.register_configs import register_configs +from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper +from brittle_star_project.environment.obs_processing import create_obs_processor +from brittle_star_project.environment.padded_obs_wrapper import compute_padding_masks +from brittle_star_project.evaluation.checkpoint import ( + load_metadata, + load_params, + metadata_to_configs, +) +from brittle_star_project.evaluation.evaluate_mjx import ( + append_checkpoint_eval_row, + build_eval_rollout_fn, + evaluate_checkpoint_mjx, +) +from brittle_star_project.trainers.PPOTrainer import reward_fn + + +def _parse_iteration(checkpoint_path: Path) -> int: + """Parse the iteration number from a checkpoint filename like `checkpoint_0042.flax`.""" + match = re.search(r"(\d+)", checkpoint_path.stem) + return int(match.group(1)) if match else -1 + + +@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") +def main(dict_cfg: DictConfig) -> None: + logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") + logger = logging.getLogger(__name__) + + cfg: BrittleStarConfig = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg) + ) + sim_cfg = cfg.simulation + eval_cfg = cfg.evaluation + + # --- Resolve the model path to find the run directory --- + model_path_str = sim_cfg.model_path + if model_path_str is None: + raise ValueError( + "simulation.model_path must point to the final_model.flax of a training run." + ) + + model_path = Path(hydra.utils.to_absolute_path(model_path_str)) + run_dir = model_path.parent + + checkpoints_dir = run_dir / "checkpoints" + if not checkpoints_dir.exists(): + raise FileNotFoundError( + f"No checkpoints/ directory found in run directory: {run_dir}\n" + "Make sure simulation.model_path points to a completed training run." + ) + + checkpoints = sorted(checkpoints_dir.glob("*.flax"), key=_parse_iteration) + if not checkpoints: + raise FileNotFoundError(f"No .flax checkpoints found in {checkpoints_dir}") + + logger.info(f"Found {len(checkpoints)} checkpoint(s) in {checkpoints_dir}") + + # --- Load sidecar metadata + reconstruct training config --- + metadata_override = ( + Path(hydra.utils.to_absolute_path(sim_cfg.metadata_path)) + if sim_cfg.metadata_path is not None + else None + ) + metadata = load_metadata(model_path, metadata_override) + training = metadata_to_configs(metadata) + + padding_masks = compute_padding_masks( + segments_per_arm=training.morphology.segments_per_arm, + reference_segments_per_arm=training.morphology.segments_per_arm, + ) + + morph_mode = training.morphology.morph_mode + + segments_per_arm = jnp.asarray( + training.morphology.segments_per_arm, + dtype=jnp.int32, + ) + + num_arms = ( + jnp.where( + segments_per_arm > 0, + 1, + 0, + ) + .sum() + .item() + ) + + match morph_mode: + case MorphMode.CENTRALIZED: + needed_copies = 1 + agent_indices = [0, 1, 2, 3, 4] + + case MorphMode.FULLY_CONNECTED | MorphMode.RING: + agent_mask = segments_per_arm > 0 + agent_indices = jnp.where(agent_mask)[0] + needed_copies = num_arms + + case MorphMode.SEGMENT: + agent_mask = segments_per_arm > 0 + agent_indices = jnp.where(agent_mask)[0] + + needed_copies = (segments_per_arm.sum() + num_arms).item() + + obs_processor = create_obs_processor( + bounds_dict=training.obs_bounds.to_bounds_dict(), + padding_masks=padding_masks, + num_arms=num_arms, + needed_copies=needed_copies, + morph_mode=morph_mode, + segments_per_arm=segments_per_arm, + agent_indices=agent_indices, + ) + + env = BrittleStarJaxEnvWrapper( + morphology=training.morphology, + arena=training.arena, + env_config=training.environment, + num_envs=1, + ) + + action_low = np.asarray(env.single_action_space.low, dtype=np.float32) + action_high = np.asarray(env.single_action_space.high, dtype=np.float32) + + sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300]) + actor = Actor(action_dim=env.single_action_space.shape[0]) + sensor.apply = jax.jit(sensor.apply) + actor.apply = jax.jit(actor.apply) + + eval_fn = build_eval_rollout_fn( + env=env, + obs_processor=obs_processor, + sensor_apply=sensor.apply, + actor_apply=actor.apply, + action_low=action_low, + action_high=action_high, + reward_fn=reward_fn, + ) + + morph_mode = training.morphology.morph_mode + + segments_per_arm = jnp.asarray( + training.morphology.segments_per_arm, + dtype=jnp.int32, + ) + + match morph_mode: + case MorphMode.CENTRALIZED: + needed_copies = 1 + + case MorphMode.FULLY_CONNECTED | MorphMode.RING: + needed_copies = jnp.where(segments_per_arm > 0, 1, 0).sum().item() + + case MorphMode.SEGMENT: + needed_copies = ( + segments_per_arm.sum() + jnp.where(segments_per_arm > 0, 1, 0).sum() + ).item() + + adj = build_adjacency( + training.morphology.segments_per_arm, + morph_mode, + ) + + sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300]) + + actor = Actor(action_dim=env.single_action_space.shape[0] // needed_copies) + + message_passer = ( + MessagePasser( + hidden_dim=300, + num_propagation_steps=4, + adj_matrix=adj, + ) + if morph_mode != MorphMode.CENTRALIZED + else None + ) + + eval_fn = build_eval_rollout_fn( + env=env, + obs_processor=obs_processor, + sensor_apply=lambda p, x: apply_per_node(sensor.apply, p, x), + actor_apply=lambda p, x: apply_per_node(actor.apply, p, x), + message_passer_apply=(None if message_passer is None else message_passer.apply), + action_low=action_low, + action_high=action_high, + reward_fn=reward_fn, + ) + seed = int(eval_cfg.eval_seed) + max_steps = int(eval_cfg.eval_max_steps) + + logger.info(f"Evaluating each checkpoint (seed={seed}, max_steps={max_steps}).") + + for checkpoint_path in checkpoints: + iteration = _parse_iteration(checkpoint_path) + try: + params = load_params(checkpoint_path) + except Exception as e: + logger.warning(f"Could not load {checkpoint_path.name}: {e}") + continue + + result = evaluate_checkpoint_mjx(eval_fn, params, seed=seed, max_steps=max_steps) + csv_path = append_checkpoint_eval_row( + run_dir, + iteration=iteration, + trained_timesteps=0, # unknown without training logs + result=result, + ) + + logger.debug( + f"checkpoint={iteration:5d} | " + f"reached={str(result.reached_target):<5} | " + f"return={result.eval_return:+8.3f} | " + f"steps={result.steps:4d} | " + f"final_dist={result.final_xy_dist:.3f}" + ) + + logger.info(f"Done. CSV at: {csv_path}") + env.close() + + +if __name__ == "__main__": + register_configs() + main() diff --git a/scripts/hpc/export_requirements.py b/scripts/hpc/export_requirements.py new file mode 100644 index 0000000..62c1298 --- /dev/null +++ b/scripts/hpc/export_requirements.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +"""Export HPC pip requirements from pyproject.toml. + +This is a LOCAL DEVELOPER UTILITY — run it on your own machine before pushing +code whenever pyproject.toml dependencies change. It reads the modules from +env/hpc/modules.txt and the full dependency list from pyproject.toml, then +writes the remainder to env/hpc/requirements.txt. +""" + +from __future__ import annotations + +import re +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] + + +def normalise(name: str) -> str: + """Normalise a PyPI package name for comparison.""" + return re.sub(r"[-_.]+", "-", name).lower() + + +def pkg_name(dep: str) -> str: + """Extract the bare package name from a PEP 508 dependency string.""" + return re.split(r"[\[=><~!;]", dep)[0].strip() + + +def main() -> None: + import tomllib + + modules_path = ROOT / "env" / "hpc" / "modules.txt" + if not modules_path.exists(): + print(f"Error: {modules_path} not found.", file=sys.stderr) + sys.exit(1) + + # Read normalized module names from base modules only + # Library modules (like PyTorch) are kept in requirements for portability + module_names = [ + normalise(line.split()[0].split("/")[0]) + for line in modules_path.read_text().splitlines() + if line.strip() and not line.startswith("#") + ] + + pyproject_path = ROOT / "pyproject.toml" + with pyproject_path.open("rb") as f: + data = tomllib.load(f) + + # Collect all dependencies, merging 'cuda' extras into base dependencies + dep_dict: dict[str, str] = {} + for dep in data.get("project", {}).get("dependencies", []): + dep_dict[normalise(pkg_name(dep))] = dep + + # Add cuda extras (takes precedence for HPC) + optional_deps = data.get("project", {}).get("optional-dependencies", {}) + for group in ["cuda"]: + for dep in optional_deps.get(group, []): + dep_dict[normalise(pkg_name(dep))] = dep + + deps = list(dep_dict.values()) + + final_deps: list[str] = [] + print("Checking dependencies against HPC module list...", file=sys.stderr) + for dep in deps: + name = normalise(pkg_name(dep)) + # Smart check: if the package name is a substring of any loaded module name + # (e.g. 'torch' in 'pytorch', 'scipy' in 'scipy-bundle') + if any(name in mod for mod in module_names): + print(f" [skip – module provider found] {dep}", file=sys.stderr) + continue + + final_deps.append(dep) + print(f" [pip] {dep}", file=sys.stderr) + + hpc_dir = ROOT / "env" / "hpc" + output_path = hpc_dir / "requirements.txt" + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text("\n".join(final_deps) + "\n") + print(f"\nWrote {len(final_deps)} requirement(s) to {output_path}", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/scripts/hpc/install.sh b/scripts/hpc/install.sh new file mode 100644 index 0000000..f88d081 --- /dev/null +++ b/scripts/hpc/install.sh @@ -0,0 +1,54 @@ +#!/bin/bash -l +# scripts/hpc/install.sh +# +# Usage (on any compute node): +# bash scripts/hpc/install.sh +# +# Batch usage: +# qsub scripts/hpc/install.sh + +#PBS -N brittlestar-install +#PBS -l walltime=00:15:00 + +set -euo pipefail + +# Preliminary status echo +echo ">>> Starting installation job $PBS_JOBID on $(hostname)..." + +if [ -n "$PBS_O_WORKDIR" ]; then + cd "$PBS_O_WORKDIR" +fi + +mkdir -p "${PBS_O_WORKDIR}/runs" + +# Mirror configs to $VSC_DATA to avoid home quota limits (3GB) +# vsc-venv manages environments relative to the requirements file +PROJ_NAME=$(basename "$PWD") +HPC_CONFIG_DIR="$VSC_DATA/$PROJ_NAME/env/hpc" +mkdir -p "$HPC_CONFIG_DIR" +cp env/hpc/*.txt "$HPC_CONFIG_DIR/" + +# Keep caches off $VSC_HOME (quota ~3 GB). +export PIP_CACHE_DIR="$VSC_SCRATCH/.cache/pip" +export UV_CACHE_DIR="$VSC_SCRATCH/.cache/uv" +mkdir -p "$PIP_CACHE_DIR" "$UV_CACHE_DIR" + +module load vsc-venv + +echo ">>> Synchronizing and activating environment (vsc-venv)..." +# cd to $VSC_DATA so vsc-venv creates its venvs/ directory there, not in $HOME. +mkdir -p "$VSC_DATA/$PROJ_NAME" +cd "$VSC_DATA/$PROJ_NAME" +set +euo pipefail +source vsc-venv --activate \ + --modules "$HPC_CONFIG_DIR/modules.txt" \ + --requirements "$HPC_CONFIG_DIR/requirements.txt" +set -euo pipefail +cd "$PBS_O_WORKDIR" + +echo '>>> Installing ipykernel...' +CLUSTER_ID="${VSC_INSTITUTE_CLUSTER:-generic}" +python -m ipykernel install --user --name="sel3_${CLUSTER_ID}" \ + --display-name "SEL3 (${CLUSTER_ID})" + +echo '>>> Done' diff --git a/scripts/hpc/train.pbs b/scripts/hpc/train.pbs new file mode 100644 index 0000000..0f95b18 --- /dev/null +++ b/scripts/hpc/train.pbs @@ -0,0 +1,75 @@ +# Production training (requires GPU at runtime): +# qsub -l gpus=1 scripts/hpc/train.pbs +# Debug/CPU training: +# qsub scripts/hpc/train.pbs + +#PBS -N brittlestar-ppo +#PBS -l nodes=1:ppn=8 +#PBS -l walltime=24:00:00 +#PBS -o runs/brittlestar-ppo.o$PBS_JOBID +#PBS -e runs/brittlestar-ppo.e$PBS_JOBID + +set -euo pipefail + +# Preliminary status echo +echo ">>> Starting training job $PBS_JOBID on $(hostname)..." + +if [ -n "$PBS_O_WORKDIR" ]; then + cd "$PBS_O_WORKDIR" +fi + +# Set up storage paths dynamically +PROJ_NAME=$(basename "$PWD") +RUN_ID="brittlestar_${PBS_JOBID}" +SCRATCH_RUNDIR="$VSC_SCRATCH/runs/$RUN_ID" +DATA_RUNDIR="$VSC_DATA/runs/$RUN_ID" +mkdir -p "$SCRATCH_RUNDIR" "$DATA_RUNDIR" runs/ + +# Keep caches off $VSC_HOME (quota ~3 GB). +export PIP_CACHE_DIR="$VSC_SCRATCH/.cache/pip" +export UV_CACHE_DIR="$VSC_SCRATCH/.cache/uv" +mkdir -p "$PIP_CACHE_DIR" "$UV_CACHE_DIR" + +module load vsc-venv + +echo ">>> Synchronizing and activating environment (vsc-venv)..." +HPC_CONFIG_DIR="$VSC_DATA/$PROJ_NAME/env/hpc" +if [ ! -d "$HPC_CONFIG_DIR" ]; then + echo "ERROR: HPC_CONFIG_DIR ($HPC_CONFIG_DIR) does not exist. Run install.sh first." + exit 1 +fi + +# cd to $VSC_DATA so vsc-venv finds its venvs/ directory there, not in $HOME. +cd "$VSC_DATA/$PROJ_NAME" +set +euo pipefail +source vsc-venv --activate \ + --modules "$HPC_CONFIG_DIR/modules.txt" \ + --requirements "$HPC_CONFIG_DIR/requirements.txt" +set -euo pipefail +cd "$PBS_O_WORKDIR" + + +echo ">>> Starting BrittleStar training..." +export MUJOCO_GL=egl +export WANDB_DIR="$SCRATCH_RUNDIR" + +export PYTHONPATH="$PBS_O_WORKDIR/src:${PYTHONPATH:-}" + +if [ -f "$VSC_DATA/$PROJ_NAME/.env" ]; then + echo ">>> Sourcing API keys from .env..." + export $(grep -v '^#' "$VSC_DATA/$PROJ_NAME/.env" | xargs) +elif [ -f "$PBS_O_WORKDIR/.env" ]; then + echo ">>> Sourcing API keys from .env..." + export $(grep -v '^#' "$PBS_O_WORKDIR/.env" | xargs) +fi + +# Run training using Hydra overrides +python scripts/train.py \ + hydra.run.dir="$SCRATCH_RUNDIR" \ + ppo=stable \ + logging=hpc + +echo ">>> Staging out results to $DATA_RUNDIR..." +cp -r "$SCRATCH_RUNDIR/." "$DATA_RUNDIR/" + +echo ">>> Done" diff --git a/scripts/plots/analyze_comparisons.py b/scripts/plots/analyze_comparisons.py new file mode 100644 index 0000000..8a66b4c --- /dev/null +++ b/scripts/plots/analyze_comparisons.py @@ -0,0 +1,418 @@ +""" +Poster Comparison Visualizations + +This script generates a Forward Velocity plot and three secondary plots (Accumulated Reward, Success +Rate, Distance Remaining). +""" + +import os + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +from plot_config import ( + BEST_PERFORMER_COLOR, + BEST_PERFORMER_TEXT, + COLORS, + LEGEND_KWARGS, + apply_style, + create_common_parser, +) + + +def load_and_preprocess_data(filepath): + """Loads CSV and prepares the metrics for plotting.""" + df = pd.read_csv(filepath) + + # Ensure success rate can be averaged numerically + if "reached_target" in df.columns: + df["reached_target"] = df["reached_target"].astype(int) + + return df + + +def _add_square_placeholders(ax, x_positions, labels): + """Adds square placeholders for images below the x-axis.""" + for x, label in zip(x_positions, labels): + # Create a roughly square rectangle in a mix of data/axes coords + # Shifted down to avoid overlapping with x-tick labels + rect = plt.Rectangle( + (x - 0.25, -0.40), + 0.5, + 0.18, + transform=ax.get_xaxis_transform(), + facecolor="#F0F0F0", + edgecolor="#A9A9A9", + linestyle="--", + zorder=1, + clip_on=False, + ) + ax.add_patch(rect) + ax.text( + x, + -0.31, + f"[ Insert {label}\nImage ]", + transform=ax.get_xaxis_transform(), + ha="center", + va="center", + fontsize=10, + color="#888888", + zorder=2, + ) + + +def plot_grouped_bar( + df, + metric_col, + ylabel, + title, + output_filename, + output_dir, + higher_is_better=True, + show_titles=False, + figsize=(12, 8), +): + """Generates and saves a highly customized grouped bar chart (grouped by Morphology).""" + grouped = ( + df.groupby(["num_active_arms", "architecture"])[metric_col] + .agg(["mean", "std"]) + .reset_index() + ) + morphologies = sorted(grouped["num_active_arms"].unique(), reverse=True) + architectures = grouped["architecture"].unique() + + fig, ax = plt.subplots(figsize=figsize) + bar_width = 0.35 + group_spacing = 1.3 + x_indices = np.arange(len(morphologies)) * group_spacing + all_bars = {} + all_means = [] + + for i, arch in enumerate(architectures): + arch_data = grouped[grouped["architecture"] == arch] + means = [ + arch_data[arch_data["num_active_arms"] == m]["mean"].values[0] + if not arch_data[arch_data["num_active_arms"] == m].empty + else 0 + for m in morphologies + ] + stds = [ + arch_data[arch_data["num_active_arms"] == m]["std"].values[0] + if not arch_data[arch_data["num_active_arms"] == m].empty + else 0 + for m in morphologies + ] + all_means.extend(means) + x_pos = x_indices + (i * bar_width) - (bar_width / 2 if len(architectures) == 2 else 0) + color = COLORS.get(arch, "#888888") + clean_label = arch.replace("_", " ").title() + bars = ax.bar( + x_pos, + means, + bar_width, + yerr=stds, + label=clean_label, + color=color, + capsize=8, + error_kw={"elinewidth": 2, "alpha": 0.7}, + ) + all_bars[arch] = (x_pos, means, stds, bars) + + for m_idx, _ in enumerate(morphologies): + m_means = {arch: all_bars[arch][1][m_idx] for arch in architectures} + best_arch = ( + max(m_means, key=m_means.get) if higher_is_better else min(m_means, key=m_means.get) + ) + best_x = all_bars[best_arch][0][m_idx] + best_y = all_bars[best_arch][1][m_idx] + best_std = all_bars[best_arch][2][m_idx] + offset = best_std + (abs(max(m_means.values())) * 0.05) if m_means.values() else 0 + ax.text( + best_x, + best_y + offset, + BEST_PERFORMER_TEXT, + ha="center", + va="bottom", + fontsize=28, + color=BEST_PERFORMER_COLOR, + ) + + # Aesthetics + ax.set_ylabel(ylabel, labelpad=15) + if show_titles: + ax.set_title(title, pad=25, fontweight="bold") + + x_ticks_pos = ( + x_indices + + bar_width # center the label in the 3 bars + + (bar_width / 2 if len(architectures) % 2 == 0 else 0) + - (bar_width / 2 if len(architectures) == 2 else 0) + ) + ax.set_xticks(x_ticks_pos) + ax.set_xticklabels([f"{m} Arms" for m in morphologies]) + ax.tick_params(axis="x") # More padding for the squares + + # X-axis at zero + ax.axhline(0, color="black", linewidth=1.5) + ax.spines["bottom"].set_visible(False) + + # Y-axis limits explicitly including 0 + if all_means: + min_val = min([*all_means, 0]) + max_val = max([*all_means, 0]) + margin = (max_val - min_val) * 0.15 if max_val != min_val else 0.1 + ax.set_ylim(min_val - margin, max_val + margin * 1.5) # Extra top margin for stars + # Format y-ticks to not have excessive decimals, include 0 + ticks = ( + [min_val, max_val] + if min_val == 0 and max_val == 0 + else sorted(list(set([min_val, 0, max_val]))) + ) + ax.set_yticks(ticks) + ax.yaxis.set_major_formatter( + plt.FuncFormatter(lambda x, _: f"{x:.2f}" if abs(x) < 10 else f"{x:.0f}") + ) + + ax.legend(**LEGEND_KWARGS, ncol=len(architectures)) + ax.set_facecolor("white") + fig.patch.set_facecolor("white") + + os.makedirs(output_dir, exist_ok=True) + base_path = os.path.join(output_dir, os.path.splitext(output_filename)[0]) + plt.savefig(f"{base_path}.png", dpi=300, bbox_inches="tight") + plt.savefig(f"{base_path}.svg", format="svg", bbox_inches="tight") + plt.close() + + +def plot_grouped_bar_alt( + df, + metric_col, + ylabel, + title, + output_filename, + output_dir, + higher_is_better=True, + show_titles=False, + figsize=(12, 8), +): + """Generates and saves a highly customized grouped bar chart (grouped by Architecture).""" + grouped = ( + df.groupby(["architecture", "num_active_arms"])[metric_col] + .agg(["mean", "std"]) + .reset_index() + ) + architectures = sorted(grouped["architecture"].unique()) + morphologies = sorted(grouped["num_active_arms"].unique(), reverse=True) + + fig, ax = plt.subplots(figsize=figsize) + bar_width = 0.8 / len(morphologies) + x_indices = np.arange(len(architectures)) + all_bars = {} + all_means = [] + + for i, m in enumerate(morphologies): + m_data = grouped[grouped["num_active_arms"] == m] + means = [ + m_data[m_data["architecture"] == arch]["mean"].values[0] + if not m_data[m_data["architecture"] == arch].empty + else 0 + for arch in architectures + ] + stds = [ + m_data[m_data["architecture"] == arch]["std"].values[0] + if not m_data[m_data["architecture"] == arch].empty + else 0 + for arch in architectures + ] + all_means.extend(means) + + # Offset bars based on morphology index + offset = (i - len(morphologies) / 2 + 0.5) * bar_width + x_pos = x_indices + offset + + # We can use a color gradient or different colors for morphologies + # For simplicity, using a colormap + color = plt.cm.viridis(i / max(1, len(morphologies) - 1)) + + bars = ax.bar( + x_pos, + means, + bar_width, + yerr=stds, + label=f"{m} Arms", + color=color, + capsize=4, + error_kw={"elinewidth": 1.5, "alpha": 0.7}, + ) + all_bars[m] = (x_pos, means, stds, bars) + + for a_idx, arch in enumerate(architectures): + a_means = {m: all_bars[m][1][a_idx] for m in morphologies} + best_m = ( + max(a_means, key=a_means.get) if higher_is_better else min(a_means, key=a_means.get) + ) + best_x = all_bars[best_m][0][a_idx] + best_y = all_bars[best_m][1][a_idx] + best_std = all_bars[best_m][2][a_idx] + offset = best_std + (abs(max(a_means.values())) * 0.05) if a_means.values() else 0 + ax.text( + best_x, + best_y + offset, + BEST_PERFORMER_TEXT, + ha="center", + va="bottom", + fontsize=20, + color=BEST_PERFORMER_COLOR, + ) + + # Aesthetics + ax.set_ylabel(ylabel, labelpad=15) + if show_titles: + ax.set_title(title + " (Alt)", pad=25, fontweight="bold") + + ax.set_xticks(x_indices) + ax.set_xticklabels([arch.replace("_", " ").title() for arch in architectures]) + ax.tick_params(axis="x", pad=25) + + # X-axis at zero + ax.axhline(0, color="black", linewidth=1.5) + ax.spines["bottom"].set_visible(False) + + if all_means: + min_val = min([*all_means, 0]) + max_val = max([*all_means, 0]) + margin = (max_val - min_val) * 0.15 if max_val != min_val else 0.1 + ax.set_ylim(min_val - margin, max_val + margin * 1.5) + ticks = ( + [min_val, max_val] + if min_val == 0 and max_val == 0 + else sorted(list(set([min_val, 0, max_val]))) + ) + ax.set_yticks(ticks) + ax.yaxis.set_major_formatter( + plt.FuncFormatter(lambda x, _: f"{x:.2f}" if abs(x) < 10 else f"{x:.0f}") + ) + + ax.legend(**LEGEND_KWARGS, ncol=len(morphologies)) + ax.set_facecolor("white") + fig.patch.set_facecolor("white") + + os.makedirs(output_dir, exist_ok=True) + base_path = os.path.join(output_dir, os.path.splitext(output_filename)[0]) + plt.savefig(f"{base_path}.png", dpi=300, bbox_inches="tight") + plt.savefig(f"{base_path}.svg", format="svg", bbox_inches="tight") + plt.close() + + +if __name__ == "__main__": + parser = create_common_parser(description="Generate comparison poster plots.") + parser.add_argument( + "input_csv", help="Path to the input CSV file containing evaluation results." + ) + args = parser.parse_args() + + INPUT_CSV = args.input_csv + OUTPUT_DIR = args.output_dir + + if not os.path.exists(INPUT_CSV): + print(f"Error: Could not find {INPUT_CSV}. Please ensure the file exists.") + else: + df = load_and_preprocess_data(INPUT_CSV) + print("Data loaded successfully. Generating poster plots...") + + apply_style(font_size=args.font_size) + kwargs = {"show_titles": args.show_titles, "figsize": (args.fig_width, args.fig_height)} + + # Velocity Conversion: m/s to cm/s + if "approx_max_velocity" in df.columns: + df["approx_max_velocity"] = df["approx_max_velocity"] * 100 + + # 1. Primary Plot: Forward Velocity + plot_grouped_bar( + df=df, + metric_col="approx_max_velocity", + ylabel="", + title="Maximal forward velocity (in cm/s)", + output_filename="poster_plot_velocity.png", + output_dir=OUTPUT_DIR, + higher_is_better=True, + **kwargs, + ) + plot_grouped_bar_alt( + df=df, + metric_col="approx_max_velocity", + ylabel="Max Forward Velocity (cm/s)", + title="Graceful Degradation: Velocity Across Morphologies", + output_filename="poster_plot_velocity_alt.png", + output_dir=OUTPUT_DIR, + higher_is_better=True, + **kwargs, + ) + + # 2. Secondary Plot: Accumulated Reward + plot_grouped_bar( + df=df, + metric_col="eval_return", + ylabel="Mean Cumulative Reward", + title="Overall Efficiency Across Morphologies", + output_filename="poster_plot_reward.png", + output_dir=OUTPUT_DIR, + higher_is_better=True, + **kwargs, + ) + plot_grouped_bar_alt( + df=df, + metric_col="eval_return", + ylabel="Mean Cumulative Reward", + title="Overall Efficiency Across Morphologies", + output_filename="poster_plot_reward_alt.png", + output_dir=OUTPUT_DIR, + higher_is_better=True, + **kwargs, + ) + + # 3. Secondary Plot: Success Rate + plot_grouped_bar( + df=df, + metric_col="reached_target", + ylabel="Success Rate (%)", + title="Target Acquisition Consistency", + output_filename="poster_plot_success_rate.png", + output_dir=OUTPUT_DIR, + higher_is_better=True, + **kwargs, + ) + plot_grouped_bar_alt( + df=df, + metric_col="reached_target", + ylabel="Success Rate (%)", + title="Target Acquisition Consistency", + output_filename="poster_plot_success_rate_alt.png", + output_dir=OUTPUT_DIR, + higher_is_better=True, + **kwargs, + ) + + # 4. Secondary Plot: Final Distance Remaining + plot_grouped_bar( + df=df, + metric_col="final_xy_dist", + ylabel="Distance to Target Remaining", + title="Navigational Accuracy (Lower is Better)", + output_filename="poster_plot_distance.png", + output_dir=OUTPUT_DIR, + higher_is_better=False, # For distance, a lower score is better + **kwargs, + ) + plot_grouped_bar_alt( + df=df, + metric_col="final_xy_dist", + ylabel="Distance to Target Remaining", + title="Navigational Accuracy (Lower is Better)", + output_filename="poster_plot_distance_alt.png", + output_dir=OUTPUT_DIR, + higher_is_better=False, + **kwargs, + ) + + print(f"All plots generated in the '{OUTPUT_DIR}/' directory.") diff --git a/scripts/plots/analyze_convergence.py b/scripts/plots/analyze_convergence.py new file mode 100644 index 0000000..bc12c79 --- /dev/null +++ b/scripts/plots/analyze_convergence.py @@ -0,0 +1,385 @@ +""" +Convergence Analysis Script for Poster Visualizations + +This script analyzes evaluation metrics from multiple training runs to determine +the convergence point of different reinforcement learning architectures. + +Workflow: +1. Loads evaluation data from the CSV files defined in FILE_MAPPING. +2. Calculates a rolling average of the reward and velocity to smooth noise. +3. Determines the convergence timestep for each metric (first time 95% of peak is reached). +4. Generates a grouped bar chart comparing convergence speed and line plots of the raw curves. + +Usage: + uv run python scripts/analysis/analyze_convergence.py + +Note: For these metrics to be valid, the evaluation CSVs must be generated with +exploration noise strictly disabled (e.g., taking the mean of the action distribution). +""" + +import logging +import os + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +from enum import Enum + +from plot_config import COLORS, apply_style, create_common_parser, LEGEND_KWARGS + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) +logger = logging.getLogger(__name__) + +# --- Globals & Configuration --- +USING_DUMMY_DATA = False +SMOOTHING_WINDOW = 3 +CONVERGENCE_THRESHOLD = 0.95 + + +class Columns(str, Enum): + # ... (rest of the file remains same, just need to update plotting functions and obtain_data) + """Column names expected in every evaluation CSV.""" + + CHECKPOINT = "checkpoint" + ARCH = "architecture" + TIMESTEPS = "trained_timesteps" + REWARD = "eval_return" + VELOCITY = "velocity" + EVAL_STEPS = "eval_steps" + FINAL_XY_DIST = "final_xy_dist" + INITIAL_XY_DIST = "initial_xy_dist" + REACHED_TARGET = "reached_target" + + +# Maps architecture display names to the path of their evaluation CSV. +# Update these paths once real evaluation data is available. +FILE_MAPPING: dict[str, str] = { + # "centralized 2 arms": "runs/dummy/dummy_centralized_2_arms.csv", + "centralized 5 arms": "runs/final-v2-centralized/checkpoint_evaluation.csv", + "decentralized fully connected": "runs/final-v2-fully-conn/checkpoint_evaluation.csv", + "decentralized ring-level": "runs/final-v2-ring/checkpoint_evaluation.csv", +} + +# Architecture profiles for dummy data generation: (max_reward, max_velocity, sigmoid_speed) +_DUMMY_PROFILES: dict[str, tuple[float, float, float]] = { + "centralized 2 arms": (300, 0.8, 1.2), + "centralized 5 arms": (450, 1.1, 1.0), + "decentralized fully connected": (500, 1.3, 0.7), + "decentralized ring-level": (480, 1.2, 0.8), + "decentralized segment-level": (520, 1.4, 0.6), +} + + +def generate_dummy_csvs(file_mapping: dict[str, str]): + """ + Generates one dummy CSV per architecture in FILE_MAPPING at their expected locations. + Skips any architecture without a defined profile. + """ + checkpoints = list(range(100, 1100, 100)) + timesteps = [cp * 10_000 for cp in checkpoints] + + for arch, path in file_mapping.items(): + if arch not in _DUMMY_PROFILES: + logger.warning(f"No dummy profile for '{arch}'. Skipping.") + continue + + m_reward, m_vel, speed = _DUMMY_PROFILES[arch] + + rows = [] + for i, ts in enumerate(timesteps): + progress = 1 / (1 + np.exp(-speed * (i - 4))) + rows.append( + { + Columns.TIMESTEPS: ts, + Columns.REWARD: m_reward * progress + np.random.normal(0, 5), + Columns.VELOCITY: m_vel * progress + np.random.normal(0, 0.02), + } + ) + + # Create parent directories if they don't exist + os.makedirs(os.path.dirname(path), exist_ok=True) + + pd.DataFrame(rows).to_csv(path, index=False) + logger.info(f"Generated dummy CSV at expected path: {path}") + + +def load_metrics(file_mapping: dict[str, str]) -> pd.DataFrame: + """ + Loads one CSV per architecture, injects the architecture name as a column, + and returns the combined DataFrame with only the required columns. + """ + required = [ + Columns.CHECKPOINT, + Columns.TIMESTEPS, + Columns.REWARD, + Columns.INITIAL_XY_DIST, + Columns.FINAL_XY_DIST, + Columns.EVAL_STEPS, + ] + dfs = [] + + for arch_name, filepath in file_mapping.items(): + if not os.path.exists(filepath): + logger.warning(f"File not found: '{filepath}'. Skipping.") + continue + + df = pd.read_csv(filepath) + + missing = [c for c in required if c not in df.columns] + if missing: + logger.warning(f"Missing columns {missing} in '{filepath}'. Skipping.") + continue + + df = df[required].copy() + df[Columns.VELOCITY] = (df[Columns.INITIAL_XY_DIST] - df[Columns.FINAL_XY_DIST]) / df[ + Columns.EVAL_STEPS + ] + df[Columns.ARCH] = arch_name + df[Columns.VELOCITY] = (df[Columns.INITIAL_XY_DIST] - df[Columns.FINAL_XY_DIST]) / df[ + Columns.EVAL_STEPS + ] + + dfs.append(df) + + return pd.concat(dfs, ignore_index=True) if dfs else pd.DataFrame() + + +def _convergence_timestep( + series: pd.Series, timesteps: pd.Series, checkpoints: pd.Series +) -> tuple[float, int, int]: + """Returns the first timestep where the smoothed series reaches 95% of its peak.""" + smoothed = series.rolling(window=SMOOTHING_WINDOW, min_periods=1).mean() + threshold = smoothed.max() * CONVERGENCE_THRESHOLD + + mask = smoothed >= threshold + first_idx = mask.idxmax() + + return timesteps.loc[first_idx], first_idx, checkpoints.loc[first_idx] + + +def analyze_convergence(df: pd.DataFrame) -> pd.DataFrame: + """ + For each architecture, determines the convergence timestep based on both + reward and velocity, returning one summary row per architecture. + """ + results = [] + + centralized_base = 0 + + for arch in df[Columns.ARCH].unique(): + arch_data = df[df[Columns.ARCH] == arch].sort_values(Columns.TIMESTEPS) + + reward_timestep, reward_checkpoint_idx, reward_checkpoint = _convergence_timestep( + arch_data[Columns.REWARD], + arch_data[Columns.TIMESTEPS], + arch_data[Columns.CHECKPOINT], + ) + + velocity_timestep, velocity_checkpoint_idx, velocity_checkpoint = _convergence_timestep( + arch_data[Columns.VELOCITY], + arch_data[Columns.TIMESTEPS], + arch_data[Columns.CHECKPOINT], + ) + + results.append( + { + "Architecture": arch, + "Reward_Convergence_Timestep": reward_timestep, + "Reward_Convergence_Checkpoint_Idx": reward_checkpoint_idx, + "Reward_Convergence_Checkpoint": reward_checkpoint, + "Velocity_Convergence_Timestep": velocity_timestep, + "Velocity_Convergence_Checkpoint_Idx": velocity_checkpoint_idx, + "Velocity_Convergence_Checkpoint": velocity_checkpoint, + } + ) + + if arch == "centralized 5 arms": + centralized_base = reward_checkpoint + else: + print(arch, "speedup:", 1 - reward_checkpoint / centralized_base) + + return pd.DataFrame(results) + + +def _add_bar_labels(bars, max_val: float): + """Annotates each bar with its value in white bold text, positioned inside.""" + for bar in bars: + width = bar.get_width() + label = f"{width / 1e6:.1f}M" if width >= 1e6 else f"{width:,.0f}" + plt.text( + width - (max_val * 0.02), + bar.get_y() + bar.get_height() / 2, + label, + ha="right", + va="center", + fontsize=11, + color="white", + fontweight="bold", + ) + + +def plot_grouped_convergence_chart( + results_df: pd.DataFrame, output_filename: str, output_dir: str, **kwargs +): + """ + Saves a grouped horizontal bar chart comparing Reward and Velocity convergence timesteps + across all architectures. + """ + sorted_df = results_df.sort_values("Reward_Convergence_Timestep", ascending=True) + architectures = sorted_df["Architecture"].tolist() + y_pos = np.arange(len(architectures)) + bar_height = 0.35 + max_val = sorted_df[ + ["Reward_Convergence_Timestep", "Velocity_Convergence_Timestep"] + ].values.max() + + fig, ax = plt.subplots(figsize=kwargs.get("figsize", (12, 8))) + + bars_reward = ax.barh( + y_pos + bar_height / 2, + sorted_df["Reward_Convergence_Timestep"], + height=bar_height, + label="Reward Convergence", + color="#1f77b4", + ) + bars_velocity = ax.barh( + y_pos - bar_height / 2, + sorted_df["Velocity_Convergence_Timestep"], + height=bar_height, + label="Velocity Convergence", + color="#ff7f0e", + ) + + title_suffix = " (DUMMY DATA)" if USING_DUMMY_DATA else "" + if kwargs.get("show_titles", True): + ax.set_title( + f"Comparison of Training Convergence Timesteps{title_suffix}", fontsize=20, pad=20 + ) + ax.set_xlabel("Timesteps to Convergence (95% of peak)", fontsize=16) + ax.set_ylabel("Architecture", fontsize=16) + ax.set_yticks(y_pos) + ax.set_yticklabels(architectures, fontsize=14) + ax.tick_params(axis="x", labelsize=14) + ax.legend(**LEGEND_KWARGS, ncol=2) + ax.set_xlim(left=0) + ax.spines["top"].set_visible(False) + ax.spines["right"].set_visible(False) + + _add_bar_labels(bars_reward, max_val) + _add_bar_labels(bars_velocity, max_val) + + plt.tight_layout() + os.makedirs(output_dir, exist_ok=True) + base_path = os.path.join(output_dir, os.path.splitext(output_filename)[0]) + plt.savefig(f"{base_path}.png", dpi=300, bbox_inches="tight") + plt.savefig(f"{base_path}.svg", format="svg", bbox_inches="tight") + plt.close() + + +def plot_metric_curves( + df: pd.DataFrame, metric_col: str, title: str, output_filename: str, output_dir: str, **kwargs +): + """ + Saves a line plot of the given metric over training timesteps for every architecture. + """ + fig, ax = plt.subplots(figsize=kwargs.get("figsize", (12, 7))) + + for arch in df[Columns.ARCH].unique(): + arch_data = df[df[Columns.ARCH] == arch].sort_values(Columns.TIMESTEPS) + color_key = arch.split()[0].upper() if isinstance(arch, str) else "UNKNOWN" + color = COLORS.get(color_key, "#888888") + ax.plot( + arch_data[Columns.TIMESTEPS], + arch_data[metric_col], + label=arch, + marker="o", + markersize=4, + alpha=0.8, + color=color, + ) + + title_suffix = " (DUMMY DATA)" if USING_DUMMY_DATA else "" + if kwargs.get("show_titles", True): + ax.set_title(f"{title}{title_suffix}", fontsize=18, pad=20) + ax.set_xlabel("Training Timesteps", fontsize=14) + ax.set_ylabel(metric_col.replace("_", " ").title(), fontsize=14) + ax.legend(**LEGEND_KWARGS, ncol=len(df[Columns.ARCH].unique())) + ax.grid(True, linestyle="--", alpha=0.6) + ax.set_xlim(left=0) + ax.set_ylim(bottom=0) + + plt.tight_layout() + os.makedirs(output_dir, exist_ok=True) + base_path = os.path.join(output_dir, os.path.splitext(output_filename)[0]) + plt.savefig(f"{base_path}.png", dpi=300, bbox_inches="tight") + plt.savefig(f"{base_path}.svg", format="svg", bbox_inches="tight") + plt.close() + + +def plot_results(df: pd.DataFrame, results: pd.DataFrame, output_dir: str, **kwargs): + """Generates and saves all analysis plots.""" + plot_grouped_convergence_chart( + results, output_filename="convergence_comparison.png", output_dir=output_dir, **kwargs + ) + plot_metric_curves( + df, + Columns.REWARD, + "Training Progress: Accumulated Reward", + "progress_reward_curves.png", + output_dir=output_dir, + **kwargs, + ) + plot_metric_curves( + df, + Columns.VELOCITY, + "Training Progress: Velocity", + "progress_velocity_curves.png", + output_dir=output_dir, + **kwargs, + ) + + +def obtain_data() -> pd.DataFrame: + """Resolves the file mapping, falling back to generated dummy CSVs if needed.""" + global USING_DUMMY_DATA + + if not any(os.path.exists(p) for p in FILE_MAPPING.values()): + logger.info("No real evaluation files found. Generating dummy CSVs at expected locations.") + generate_dummy_csvs(FILE_MAPPING) + USING_DUMMY_DATA = True + + return load_metrics(FILE_MAPPING) + + +def run_analysis(output_dir: str, **kwargs): + """Orchestrates data loading, convergence analysis, and plot generation.""" + df = obtain_data() + if df.empty: + logger.error("No data found to analyze.") + return + + results = analyze_convergence(df) + print( + results[ + ["Architecture", "Reward_Convergence_Checkpoint_Idx", "Reward_Convergence_Checkpoint"] + ] + ) + + plot_results(df, results, output_dir, **kwargs) + logger.info("Analysis complete. Plots saved to disk.") + + +if __name__ == "__main__": + parser = create_common_parser(description="Analyze training convergence.") + args = parser.parse_args() + + apply_style(font_size=args.font_size) + run_analysis( + output_dir=args.output_dir, + show_titles=args.show_titles, + figsize=(args.fig_width, args.fig_height), + ) diff --git a/scripts/plots/plot_config.py b/scripts/plots/plot_config.py new file mode 100644 index 0000000..48f9106 --- /dev/null +++ b/scripts/plots/plot_config.py @@ -0,0 +1,75 @@ +import argparse +import matplotlib.pyplot as plt + +# Shared Color Palette (Colorblind friendly, high contrast) +# Matches poster design +COLORS = { + "CENTRALIZED": "#0D567C", # Blue + "FULLY_CONNECTED": "#8C0E0F", # Reddish + "RING": "#FCB305", # Pale Yellow +} + + +def apply_style(font_size=36): + """ + Applies the shared typography and aesthetic settings to Matplotlib. + """ + plt.rcParams.update( + { + "font.size": font_size, + "axes.labelsize": font_size, + "axes.titlesize": font_size, + "xtick.labelsize": font_size, + "ytick.labelsize": font_size, + "legend.fontsize": font_size, + "axes.linewidth": 2, + "axes.spines.top": False, + "axes.spines.right": False, + "axes.spines.left": False, + "figure.facecolor": "white", + "axes.facecolor": "white", + "savefig.bbox": "tight", + "savefig.dpi": 300, + } + ) + + +# Star marker for best performer +BEST_PERFORMER_TEXT = "★" +BEST_PERFORMER_MARKER = "*" +BEST_PERFORMER_COLOR = "#D4AF37" # Gold + +# Centralized Legend Configuration +LEGEND_KWARGS = { + "loc": "upper center", + "bbox_to_anchor": (0.5, -0.12), + "frameon": False, +} + + +def create_common_parser(description: str) -> argparse.ArgumentParser: + """ + Creates an argparse parser with common plotting arguments. + """ + parser = argparse.ArgumentParser(description=description) + parser.add_argument( + "--output_dir", + "-o", + default="runs/evaluation/plots", + help="Directory to save the generated plots.", + ) + parser.add_argument( + "--show_titles", + action="store_true", + help="Include titles in the plots. Default is False for easier poster integration.", + ) + parser.add_argument( + "--font_size", type=int, default=28, help="Base font size in points. Default is 28." + ) + parser.add_argument( + "--fig_width", type=float, default=12.0, help="Figure width in inches. Default is 12.0." + ) + parser.add_argument( + "--fig_height", type=float, default=8.0, help="Figure height in inches. Default is 8.0." + ) + return parser diff --git a/scripts/poster_visualisations/render_poster_videos.py b/scripts/poster_visualisations/render_poster_videos.py new file mode 100644 index 0000000..eca60fc --- /dev/null +++ b/scripts/poster_visualisations/render_poster_videos.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +import argparse +from pathlib import Path + +from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs +from brittle_star_project.evaluation.eval_env_builder import build_eval_env +from brittle_star_project.evaluation.video import record_episode_multi_camera + +_ARCH_DIR_MAP = { + "CENTRALIZED": "centralized", + "FULLY_CONNECTED": "fully-connected", + "RING": "ring", + "SEGMENT": "segment", +} + +ROBOT_COLOR_MAP = { + "CENTRALIZED": "#0D567C", # Blue + "FULLY_CONNECTED": "#8C0E0F", # Reddish + "RING": "#FCB304", # Pale Yellow +} + + +def _arch_dir(name: str) -> str: + return _ARCH_DIR_MAP.get(name, name.lower()) + + +def _resolve_overrides(overrides: list[str], count: int) -> list[str | None]: + if not overrides: + return [None] * count + if len(overrides) == 1 and count > 1: + return overrides * count + if len(overrides) != count: + raise ValueError("morphology overrides must match the number of models") + return overrides + + +def main() -> None: + parser = argparse.ArgumentParser(description="Render top-down and follow videos for poster.") + parser.add_argument("models", nargs="+", help="Paths to .flax checkpoints") + parser.add_argument( + "--morphology-override", + action="append", + default=[], + help="Override morphology YAML path (repeat to match models)", + ) + parser.add_argument("--output-root", default="vids/poster") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--max-steps", type=int, default=5000) + parser.add_argument("--topdown-camera", type=int, default=0) + parser.add_argument("--follow-camera", type=int, default=1) + parser.add_argument("--topdown-camera-x", type=float, default=-3.0) + parser.add_argument("--topdown-camera-y", type=float, default=0.0) + parser.add_argument("--topdown-camera-z", type=float, default=4.5) + parser.add_argument("--topdown-camera-fovy", type=float, default=None) + parser.add_argument("--target-x", type=float, default=-6.0) + parser.add_argument("--target-y", type=float, default=0.0) + parser.add_argument("--width", type=int, default=2160) + parser.add_argument("--height", type=int, default=960) + parser.add_argument("--fps", type=int, default=60) + parser.add_argument( + "--robot-color", + default="#2B4162", + help="Hex color for the brittle star robot", + ) + args = parser.parse_args() + + if (args.target_x is None) != (args.target_y is None): + raise ValueError("target-x and target-y must be provided together") + + target_xy = None + if args.target_x is not None: + target_xy = (float(args.target_x), float(args.target_y)) + + camera_fovy = None + if args.topdown_camera_fovy is not None: + camera_fovy = {args.topdown_camera: float(args.topdown_camera_fovy)} + + camera_x = None + if args.topdown_camera_x is not None: + camera_x = {args.topdown_camera: float(args.topdown_camera_x)} + + camera_y = None + if args.topdown_camera_y is not None: + camera_y = {args.topdown_camera: float(args.topdown_camera_y)} + + camera_z = None + if args.topdown_camera_z is not None: + camera_z = {args.topdown_camera: float(args.topdown_camera_z)} + + camera_xyz = (camera_x, camera_y, camera_z) + + overrides = _resolve_overrides(args.morphology_override, len(args.models)) + output_root = Path(args.output_root) + + for model_path_str, override in zip(args.models, overrides): + model_path = Path(model_path_str) + metadata = load_metadata(model_path, None) + training = metadata_to_configs(metadata) + + bundle = build_eval_env( + model_path=model_path, + training=training, + metadata=metadata, + morphology_override_path=override, + ) + + arch_dir = _arch_dir(bundle.architecture) + arms_dir = f"{bundle.num_active_arms}arms" + out_dir = output_root / arms_dir / arch_dir + out_dir.mkdir(parents=True, exist_ok=True) + + output_paths = { + args.topdown_camera: out_dir / "topdown.mp4", + args.follow_camera: out_dir / "follow.mp4", + } + + print(bundle.architecture) + color = ROBOT_COLOR_MAP.get(bundle.architecture, args.robot_color) + + result = record_episode_multi_camera( + env=bundle.env, + policy=bundle.policy, + seed=args.seed, + max_steps=args.max_steps, + action_low=bundle.action_low, + action_high=bundle.action_high, + action_mask=bundle.action_mask, + output_paths=output_paths, + camera_ids=[args.topdown_camera, args.follow_camera], + camera_fovy=camera_fovy, + camera_xyz=camera_xyz, + target_xy=target_xy, + robot_color=color, + width=args.width, + height=args.height, + fps=args.fps, + ) + + final_dist = "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}" + print( + f"{arms_dir}/{arch_dir}: return={result.return_:.6f}, len={result.length}, " + f"target_reached={result.reached_target}, final_xy_dist={final_dist}" + ) + + bundle.env.close() + + +if __name__ == "__main__": + main() diff --git a/scripts/poster_visualisations/render_poster_videos.sh b/scripts/poster_visualisations/render_poster_videos.sh new file mode 100755 index 0000000..10648f3 --- /dev/null +++ b/scripts/poster_visualisations/render_poster_videos.sh @@ -0,0 +1,10 @@ +#!/usr/bin/env bash + +# Multi-camera renders per model -> vids/poster/{arms}arms/{arch}/topdown.mp4 + follow.mp4 + +path=$1 + +uv run scripts/poster_visualisations/render_poster_videos.py \ + "$path"/centralized.flax \ + "$path"/fully-connected.flax \ + "$path"/ring.flax \ \ No newline at end of file diff --git a/scripts/poster_visualisations/render_static_path_image.py b/scripts/poster_visualisations/render_static_path_image.py new file mode 100644 index 0000000..8229f52 --- /dev/null +++ b/scripts/poster_visualisations/render_static_path_image.py @@ -0,0 +1,241 @@ +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np + +from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs +from brittle_star_project.evaluation.eval_env_builder import build_eval_env +from brittle_star_project.evaluation.rollout import ( + _get_observations, + _maybe_clip_action, + _target_reached, +) +from brittle_star_project.evaluation.video import ( + _apply_camera_overrides, + _ensure_offscreen_size, + hex_to_rgba, +) + +ROBOT_COLOR_MAP = { + "CENTRALIZED": "#0D567C", # Blue + "FULLY_CONNECTED": "#8C0E0F", # Reddish + "RING": "#FCB304", # Pale Yellow +} + + +def _enum_value(enum_obj, *names: str) -> int: + for name in names: + if hasattr(enum_obj, name): + return int(getattr(enum_obj, name)) + raise AttributeError(f"Could not find any of {names!r} on {enum_obj!r}") + + +def _append_sphere(scene, mujoco, center: np.ndarray, radius: float, rgba: np.ndarray) -> None: + geom = scene.geoms[scene.ngeom] + mujoco.mjv_initGeom( + geom, + mujoco.mjtGeom.mjGEOM_SPHERE, + np.asarray([radius, 0.0, 0.0], dtype=np.float32), + center, + np.eye(3, dtype=np.float32).reshape(-1), + rgba, + ) + scene.ngeom += 1 + + +def main() -> None: + parser = argparse.ArgumentParser(description="Render a static path image from a rollout.") + parser.add_argument("model", help="Path to .flax checkpoint") + parser.add_argument("--morphology-override", default=None) + parser.add_argument("--output-path", required=True) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--max-steps", type=int, default=5000) + parser.add_argument("--body-name", default="BrittleStarMorphology/central_disk") + parser.add_argument("--camera-id", type=int, default=0) + parser.add_argument("--camera-x", type=float, default=-3.0) + parser.add_argument("--camera-y", type=float, default=0.0) + parser.add_argument("--camera-z", type=float, default=4.5) + parser.add_argument("--camera-fovy", type=float, default=None) + parser.add_argument("--target-x", type=float, default=-6.0) + parser.add_argument("--target-y", type=float, default=0.0) + parser.add_argument("--width", type=int, default=2160) + parser.add_argument("--height", type=int, default=960) + parser.add_argument("--frame-stride", type=int, default=15) + parser.add_argument( + "--path-color", default="#FA9F42" + ) # ring = #888888, centralized = #2B4162, fully connected = FA9F42 + parser.add_argument( + "--robot-color", + default="#FA9F42", + help="Hex color for brittle star robot (e.g. #ff0000)", + ) + args = parser.parse_args() + + model_path = Path(args.model) + metadata = load_metadata(model_path, None) + training = metadata_to_configs(metadata) + + bundle = build_eval_env( + model_path=model_path, + training=training, + metadata=metadata, + morphology_override_path=args.morphology_override, + ) + + if (args.target_x is None) != (args.target_y is None): + raise ValueError("target-x and target-y must be provided together") + + target_xy = None + if args.target_x is not None: + target_xy = (float(args.target_x), float(args.target_y)) + + try: + import imageio + import mujoco + except ImportError as e: + raise ImportError( + "Static image rendering requires 'mujoco' and 'imageio'. " + "Please install the evaluation dependencies: `uv pip install .[evaluation]`" + ) from e + + reset_kwargs = {} + if target_xy is not None: + reset_kwargs["target_position"] = (target_xy[0], target_xy[1], 0.0) + + state = bundle.env.reset(seed=args.seed, **reset_kwargs) + model = state.mj_model + data = state.mj_data + + _apply_camera_overrides( + model, + camera_fovy={args.camera_id: float(args.camera_fovy)} + if args.camera_fovy is not None + else None, + camera_xyz=( + {args.camera_id: float(args.camera_x)} if args.camera_x is not None else None, + {args.camera_id: float(args.camera_y)} if args.camera_y is not None else None, + {args.camera_id: float(args.camera_z)} if args.camera_z is not None else None, + ), + ) + _ensure_offscreen_size(model, args.width, args.height) + + body_id = mujoco.mj_name2id( + model, mujoco.mjtObj.mjOBJ_BODY, "BrittleStarMorphology/central_disk" + ) + + robot_rgba = hex_to_rgba(ROBOT_COLOR_MAP.get(bundle.architecture, args.robot_color), 1.0) + + # Optionally override robot color by recoloring geoms belonging to the robot's body subtree. + if args.robot_color is not None: + # Collect body IDs in the subtree rooted at `body_id` by walking parent links. + nbody = int(model.nbody) + body_parent = model.body_parentid + robot_body_ids = set([int(body_id)]) + for i in range(1, nbody): + cur = int(i) + # walk up until root (0) or until we hit the robot root + while cur not in (-1, 0, int(body_id)): + cur = int(body_parent[cur]) + if cur == int(body_id): + robot_body_ids.add(i) + + # Recolor geoms whose body id is in the robot subtree + for g in range(int(model.ngeom)): + if int(model.geom_bodyid[g]) in robot_body_ids: + model.geom_rgba[g][:] = robot_rgba + + # Optionally override robot color by recoloring geoms belonging to the robot's body subtree. + if args.robot_color is not None: + # Collect body IDs in the subtree rooted at `body_id` by walking parent links. + nbody = int(model.nbody) + body_parent = model.body_parentid + robot_body_ids = set([int(body_id)]) + for i in range(1, nbody): + cur = int(i) + # walk up until root (0) or until we hit the robot root + while cur not in (-1, 0, int(body_id)): + cur = int(body_parent[cur]) + if cur == int(body_id): + robot_body_ids.add(i) + + # Recolor geoms whose body id is in the robot subtree + for g in range(int(model.ngeom)): + if int(model.geom_bodyid[g]) in robot_body_ids: + model.geom_rgba[g][:] = robot_rgba + + positions = [] + observations = _get_observations(state) + + for _ in range(int(args.max_steps)): + positions.append(np.asarray(data.xpos[body_id], dtype=np.float32)) + + obs_dict = observations or {} + action = bundle.policy.act(observations=obs_dict) + if bundle.action_mask is not None: + action = action[bundle.action_mask] + action = _maybe_clip_action(action, bundle.action_low, bundle.action_high) + + state = bundle.env.step(state=state, action=action) + data = state.mj_data + observations = _get_observations(state) + + if _target_reached(state=state): + break + + positions_arr = np.vstack(positions) + if len(positions_arr) < 2: + raise ValueError("Need at least two rollout positions to render a path") + + path_points = positions_arr.copy() + path_points[:, 2] -= 0.02 + + path_step = max(1, int(args.frame_stride)) + path_points_visible = path_points[::path_step] + path_rgba = hex_to_rgba(ROBOT_COLOR_MAP.get(bundle.architecture, args.path_color), 0.92) + + ctx = mujoco.GLContext(args.width, args.height) + ctx.make_current() + try: + catmask = _enum_value(mujoco.mjtCatBit, "mjCAT_ALL") + camera_type = _enum_value(mujoco.mjtCamera, "mjCAMERA_FIXED") + font_scale = _enum_value(mujoco.mjtFontScale, "mjFONTSCALE_100") + + maxgeom = int(model.ngeom + len(path_points_visible) + 8) + scene = mujoco.MjvScene(model, maxgeom=maxgeom) + option = mujoco.MjvOption() + perturb = mujoco.MjvPerturb() + camera = mujoco.MjvCamera() + mujoco.mjv_defaultOption(option) + mujoco.mjv_defaultPerturb(perturb) + mujoco.mjv_defaultCamera(camera) + camera.type = camera_type + camera.fixedcamid = int(args.camera_id) + if hasattr(camera, "trackbodyid"): + camera.trackbodyid = -1 + + context = mujoco.MjrContext(model, font_scale) + viewport = mujoco.MjrRect(0, 0, args.width, args.height) + + mujoco.mjv_updateScene(model, data, option, perturb, camera, catmask, scene) + + for idx, path_point in enumerate(path_points_visible): + path_rgba[3] = 0.10 + 0.70 * (idx / max(len(path_points_visible) - 1, 1)) + _append_sphere(scene, mujoco, path_point, 0.03, path_rgba) + + rgb = np.empty((args.height, args.width, 3), dtype=np.uint8) + depth = np.empty((args.height, args.width), dtype=np.float32) + mujoco.mjr_render(viewport, scene, context) + mujoco.mjr_readPixels(rgb, depth, viewport, context) + imageio.imwrite(args.output_path, np.flipud(rgb)) + + context.free() + finally: + ctx.free() + + bundle.env.close() + + +if __name__ == "__main__": + main() diff --git a/scripts/poster_visualisations/render_static_path_image.sh b/scripts/poster_visualisations/render_static_path_image.sh new file mode 100755 index 0000000..ab21ca0 --- /dev/null +++ b/scripts/poster_visualisations/render_static_path_image.sh @@ -0,0 +1,10 @@ +#!/usr/bin/env bash + +# Static path image + optional ghost render + +path=$1 # path to .flax model with metadata.yaml alongside it + +uv run scripts/poster_visualisations/render_static_path_image.py \ + "$path" \ + --output-path vids/poster/5arms/centralized/path.png \ + --ghost-overlay diff --git a/scripts/simulate.py b/scripts/simulate.py new file mode 100644 index 0000000..db8f70d --- /dev/null +++ b/scripts/simulate.py @@ -0,0 +1,180 @@ +"""Simulate a trained policy in the MuJoCo viewer. + +Automatically extracts the training configuration (morphology, environment, etc.) +from the sidecar metadata YAML file to ensure simulation perfectly matches training. +Override simulation settings via CLI, e.g.: + uv run scripts/simulate.py \ + simulation.morphology_override=configs/morphology/3_arms.yaml \ + simulation.model_path=runs/.../final_model.flax +""" + +from __future__ import annotations + +from pathlib import Path + +import hydra +from omegaconf import DictConfig, OmegaConf + + +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.configs.register_configs import register_configs + +from brittle_star_project.evaluation.checkpoint import load_metadata, metadata_to_configs +from brittle_star_project.evaluation.eval_env_builder import build_eval_env +from brittle_star_project.evaluation.rollout import rollout_headless, rollout_viewer +from brittle_star_project.evaluation.video import ( + record_episode, + create_evaluation_dir, + save_evaluation_metadata, +) + + +@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") +def main(dict_cfg: DictConfig) -> None: + # 1. Hydra composes ONLY SimulationSettings + cfg = OmegaConf.to_object(OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg)) + sim_cfg = cfg.simulation + + model_path_str = sim_cfg.model_path + if model_path_str is None: + raise ValueError( + "simulation.model_path must be set to a .flax checkpoint (e.g. final_model.flax)" + ) + + model_path = Path(hydra.utils.to_absolute_path(model_path_str)) + if model_path.suffix != ".flax": + raise ValueError(f"Expected a '.flax' checkpoint, got '{model_path.name}'.") + + # 2. Discover + load sidecar metadata YAML + metadata_override = None + if sim_cfg.metadata_path is not None: + metadata_override = Path(hydra.utils.to_absolute_path(sim_cfg.metadata_path)) + + metadata = load_metadata(model_path, metadata_override) + + # 3. Reconstruct typed configs from metadata + training = metadata_to_configs(metadata) + + seed = int(cfg.experiment.seed) + + # 4-7. Build evaluation environment and policy + override_path = None + if sim_cfg.morphology_override is not None: + override_path = Path(hydra.utils.to_absolute_path(sim_cfg.morphology_override)) + + bundle = build_eval_env( + model_path=model_path, + training=training, + metadata=metadata, + morphology_override_path=override_path, + ) + + env = bundle.env + policy = bundle.policy + action_low = bundle.action_low + action_high = bundle.action_high + action_mask = bundle.action_mask + + state0 = env.reset(seed=seed) + + # 8. Run simulation + headless = bool(sim_cfg.headless) + max_steps = sim_cfg.max_steps + + if sim_cfg.record_video: + if max_steps is None: + raise ValueError("simulation.max_steps is required when simulation.record_video=true") + + max_steps_i = int(max_steps) + if max_steps_i <= 0: + raise ValueError("simulation.max_steps must be > 0") + + if sim_cfg.video_output_path is None: + eval_dir = create_evaluation_dir(model_path) + output_path = eval_dir / "simulation.mp4" + else: + output_path = Path(hydra.utils.to_absolute_path(sim_cfg.video_output_path)) + eval_dir = output_path.parent + eval_dir.mkdir(parents=True, exist_ok=True) + + result = record_episode( + env=env, + policy=policy, + seed=seed, + max_steps=max_steps_i, + action_low=action_low, + action_high=action_high, + action_mask=action_mask, + output_path=output_path, + camera_id=sim_cfg.camera_id, + width=sim_cfg.video_width, + height=sim_cfg.video_height, + fps=sim_cfg.video_fps, + ) + + save_evaluation_metadata( + eval_dir=eval_dir, + morphology_override_path=sim_cfg.morphology_override, + seed=seed, + max_steps=max_steps_i, + result=result, + ) + final_dist_str = "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}" + print(f"Video saved to {output_path}") + print( + "episode done: " + f"return={result.return_:.6f}, len={result.length}, " + f"target_reached={result.reached_target}, final_xy_dist={final_dist_str}" + ) + elif headless: + if max_steps is None: + raise ValueError("simulation.max_steps is required when simulation.headless=true") + + max_steps_i = int(max_steps) + if max_steps_i <= 0: + raise ValueError("simulation.max_steps must be > 0") + + result = rollout_headless( + env=env, + policy=policy, + seed=seed, + max_steps=max_steps_i, + action_low=action_low, + action_high=action_high, + action_mask=action_mask, + ) + final_dist_str = "n/a" if result.final_xy_dist is None else f"{result.final_xy_dist:.3f}" + print( + "episode done: " + f"return={result.return_:.6f}, len={result.length}, " + f"target_reached={result.reached_target}, final_xy_dist={final_dist_str}" + ) + else: + max_steps_val = None + if max_steps is not None: + max_steps_i = int(max_steps) + if max_steps_i <= 0: + raise ValueError("simulation.max_steps must be > 0") + max_steps_val = max_steps_i + + model_dt = float(state0.mj_model.opt.timestep) + control_dt = model_dt * float(training.environment.num_physics_steps_per_control_step) + + rollout_viewer( + env=env, + policy=policy, + seed=seed, + state=state0, + control_dt=control_dt, + max_steps=max_steps_val, + action_low=action_low, + action_high=action_high, + action_mask=action_mask, + ) + + env.close() + + +if __name__ == "__main__": + register_configs() + main() diff --git a/scripts/simulate.sh b/scripts/simulate.sh new file mode 100755 index 0000000..1e6901d --- /dev/null +++ b/scripts/simulate.sh @@ -0,0 +1,9 @@ +#!/usr/bin/env bash + +path=$1 + +uv run simulate.py \ + simulation.model_path="$path"/final_model.flax \ + simulation.record_video=True \ + simulation.video_output_path=../vids/simulation.mp4 \ + simulation.max_steps=10000 diff --git a/scripts/tools/download_wandb_project.py b/scripts/tools/download_wandb_project.py new file mode 100644 index 0000000..bf2fb11 --- /dev/null +++ b/scripts/tools/download_wandb_project.py @@ -0,0 +1,142 @@ +from pathlib import Path +from concurrent.futures import ThreadPoolExecutor, as_completed +import threading +import wandb +import argparse + +# tune these depending on network / W&B limits +MAX_RUN_WORKERS = 8 +MAX_FILE_WORKERS = 16 +MAX_ARTIFACT_WORKERS = 8 + +api = wandb.Api() + +print_lock = threading.Lock() + + +def safe_print(*args, **kwargs): + with print_lock: + print(*args, **kwargs) + + +def download_file(file, run_dir): + target = run_dir / file.name + + try: + # skip existing files + if target.exists(): + return f"SKIP FILE {target}" + + target.parent.mkdir(parents=True, exist_ok=True) + + file.download(root=run_dir, replace=False) + + return f"DONE FILE {target}" + + except Exception as e: + return f"FAIL FILE {target}: {e}" + + +def sanitize_artifact_name(name: str): + return name.replace(":", "_") + + +def download_artifact(artifact, artifact_root): + try: + artifact_name = sanitize_artifact_name(artifact.name) + artifact_dir = artifact_root / artifact_name + + if artifact_dir.exists() and any(artifact_dir.iterdir()): + return f"SKIP ARTIFACT {artifact.name}" + + artifact_dir.mkdir(parents=True, exist_ok=True) + + artifact.download(root=artifact_dir) + + return f"DONE ARTIFACT {artifact.name}" + + except Exception as e: + return f"FAIL ARTIFACT {artifact.name}: {e}" + + +def download_run(run, root): + run_dir = root / f"{run.name}" + run_dir.mkdir(parents=True, exist_ok=True) + + safe_print(f"\n=== {run.name} ({run.id}) ===") + + # ------------------------- + # Download regular run files + # ------------------------- + files = list(run.files()) + + with ThreadPoolExecutor(max_workers=MAX_FILE_WORKERS) as executor: + futures = [executor.submit(download_file, file, run_dir) for file in files] + + for future in as_completed(futures): + safe_print(future.result()) + + # ------------------------- + # Download logged artifacts + # ------------------------- + artifact_root = run_dir / "artifacts" + + try: + artifacts = list(run.logged_artifacts()) + safe_print(f"Found {len(artifacts)} artifacts for {run.name}") + + with ThreadPoolExecutor(max_workers=MAX_ARTIFACT_WORKERS) as executor: + futures = [ + executor.submit(download_artifact, artifact, artifact_root) + for artifact in artifacts + ] + + for future in as_completed(futures): + safe_print(future.result()) + + except Exception as e: + safe_print(f"Artifact download failed for {run.name}: {e}") + + # ------------------------- + # OPTIONAL: download used/input artifacts + # ------------------------- + # try: + # used_artifacts = list(run.used_artifacts()) + # used_root = run_dir / "used_artifacts" + # + # for artifact in used_artifacts: + # download_artifact(artifact, used_root) + # except Exception as e: + # safe_print(f"Used artifact download failed: {e}") + + safe_print(f"Finished {run.name}") + + +def main(entity: str, project: str, root: Path): + root.mkdir(exist_ok=True) + + runs = list(api.runs(f"{entity}/{project}")) + + safe_print(f"Found {len(runs)} runs") + + with ThreadPoolExecutor(max_workers=MAX_RUN_WORKERS) as executor: + futures = [executor.submit(download_run, run, root) for run in runs] + + for future in as_completed(futures): + try: + future.result() + except Exception as e: + safe_print("RUN FAILED:", e) + + safe_print("\nAll downloads complete.") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--entity", type=str, default="SEL3-2026-Groep-4") + parser.add_argument("--project", type=str, required=True) + parser.add_argument("--root", type=str, default="runs") + args = parser.parse_args() + + root = Path(args.root) + main(entity=args.entity, project=args.project, root=root) diff --git a/scripts/tools/dump_mjcf.py b/scripts/tools/dump_mjcf.py new file mode 100644 index 0000000..ae582c7 --- /dev/null +++ b/scripts/tools/dump_mjcf.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +""" +Dump MJCF XML for a brittle-star morphology using the project's Hydra configs. + +Usage examples: + + # Use a named morphology config from configs/morphology (Hydra style) + uv run python scripts/analysis/dump_mjcf.py morphology=3_arms + + # Use a morphology override YAML (same key as simulation.morphology_override) + uv run python scripts/analysis/dump_mjcf.py \ + simulation.morphology_override=configs/morphology/3_arms.yaml + +Output path: + Provide `dump_out=path/to/file.xml` on the command line, otherwise writes `morphology.xml` in + current directory or `runs/morphologies/.xml`. +""" + +from __future__ import annotations + +import dataclasses +import logging +import sys +from pathlib import Path +from typing import Any, Optional + +import hydra +import yaml +from omegaconf import DictConfig, OmegaConf + +from brittle_star_project.configs.register_configs import register_configs +from brittle_star_project.environment.env_config import MorphologyConfig +from brittle_star_project.environment.factory import BrittleStarEnvFactory + +logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") +logger = logging.getLogger(__name__) + + +def extract_xml_string(obj: Any) -> Optional[str]: + """ + Attempts to serialize the morphology object to an XML string by checking + common dm_control and internal API methods. + """ + serialization_methods = [ + "to_xml_string", + "to_xml", + "to_string", + "to_mjcf", + "to_mjcf_string", + "get_mjcf", + "get_mjcf_str", + "export_to_xml_string", + ] + + # If the object itself has an 'mjcf' attribute, try to serialize that instead + target_obj = getattr(obj, "mjcf", obj) + + for method_name in serialization_methods: + method = getattr(target_obj, method_name, None) + if callable(method): + try: + xml_data = method() + # Safely handle both string and byte responses + if isinstance(xml_data, str): + return xml_data + elif isinstance(xml_data, bytes): + return xml_data.decode("utf-8") + except Exception as e: + logger.debug(f"Method {method_name}() failed during serialization: {e}") + + return None + + +def resolve_output_path(cfg: DictConfig) -> Path: + """Determines the appropriate output path for the MJCF XML.""" + dump_out = cfg.get("dump_out", None) + if dump_out is not None: + return Path(hydra.utils.to_absolute_path(str(dump_out))) + + morph_name = "morphology" + for arg in sys.argv[1:]: + if arg.startswith("morphology="): + morph_name = arg.split("=", 1)[1] + break + + default_out = ( + f"runs/morphologies/{morph_name}.xml" if morph_name != "morphology" else "morphology.xml" + ) + return Path(hydra.utils.to_absolute_path(default_out)) + + +@hydra.main(config_path="../../configs", config_name="main_config", version_base="1.3") +def main(cfg: DictConfig) -> None: + """Main entry point to construct the morphology and dump its XML.""" + logger.info("Initializing morphology construction...") + + # Extract morphology config safely using dict `.get()` to avoid OmegaConf AttributeErrors + simulation_cfg = cfg.get("simulation", cfg) + override_path = simulation_cfg.get("morphology_override", None) + + if override_path: + logger.info(f"Using morphology override: {override_path}") + with open(hydra.utils.to_absolute_path(override_path), "r") as f: + data = yaml.safe_load(f) or {} + morph_cfg = MorphologyConfig(**data) + else: + # Fallback to default simulation morphology, or an empty base config + morph_node = simulation_cfg.get("morphology", cfg.get("morphology", None)) + + if morph_node is not None: + # Convert OmegaConf node to dict and instantiate MorphologyConfig. + # This ensures any missing keys gracefully fall back to the dataclass defaults. + morph_dict = OmegaConf.to_container(morph_node, resolve=True) + if isinstance(morph_dict, dict): + # Filter to avoid unexpected kwargs if the dataclass is strictly defined + if dataclasses.is_dataclass(MorphologyConfig): + valid_keys = {f.name for f in dataclasses.fields(MorphologyConfig)} + morph_dict = {k: v for k, v in morph_dict.items() if k in valid_keys} + morph_cfg = MorphologyConfig(**morph_dict) + else: + morph_cfg = MorphologyConfig() + else: + morph_cfg = MorphologyConfig() + + morphology = BrittleStarEnvFactory.create_morphology(morph_cfg) + + xml_text = extract_xml_string(morphology) + if not xml_text: + raise RuntimeError("Failed to serialize morphology to MJCF/XML. ") + + out_path = resolve_output_path(cfg) + out_path.parent.mkdir(parents=True, exist_ok=True) + with out_path.open("w", encoding="utf-8") as f: + f.write(xml_text) + + logger.info(f"Successfully exported MJCF XML to: {out_path}") + + +if __name__ == "__main__": + register_configs() + main() diff --git a/scripts/tools/extract_observation_bounds.py b/scripts/tools/extract_observation_bounds.py new file mode 100644 index 0000000..374babc --- /dev/null +++ b/scripts/tools/extract_observation_bounds.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python3 +"""Empirically extract observation bounds (focused on joint velocities). + +This script creates a MuJoCo environment using the project's factory and +randomly samples actions to discover observed maxima for selected +observation keys (joint_velocity, joint_position, joint_actuator_force). + +Usage: + python scripts/extract_observation_bounds.py \ + --morphology configs/morphology/3_arms.yaml --num-steps 5000 --seed 42 + +If `--morphology` is omitted the default `MorphologyConfig()` is used. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path +import yaml +import numpy as np + +from brittle_star_project import BrittleStarEnvFactory, BrittleStarEnv, Backend +from brittle_star_project.environment.env_config import ( + MorphologyConfig, + ArenaConfig, + EnvConfig, +) + + +def load_morphology(path: str | None) -> MorphologyConfig: + if path is None: + return MorphologyConfig() + p = Path(path) + if not p.exists(): + raise FileNotFoundError(f"Morphology file not found: {p}") + with open(p, "r") as f: + data = yaml.safe_load(f) or {} + return MorphologyConfig(**data) + + +def _extract_observations(state): + # Under different backends the returned state may be a dict or an object + obs = getattr(state, "observations", None) + if obs is None and isinstance(state, dict): + obs = state.get("observations", state) + return obs + + +def find_empirical_bounds( + morph_cfg: MorphologyConfig, + arena_cfg: ArenaConfig, + env_cfg: EnvConfig, + num_steps: int = 5000, + seed: int = 42, +) -> None: + factory = BrittleStarEnvFactory() + raw_env = factory.create_environment(Backend.MJC, morph_cfg, arena_cfg, env_cfg) + env = BrittleStarEnv(raw_env, backend=Backend.MJC, config=env_cfg, morphology_config=morph_cfg) + + # Initial reset + state = env.reset(seed=seed) + + # Determine action bounds + action_space = getattr(raw_env, "action_space", None) + if action_space is None: + raise RuntimeError("Environment missing `action_space`; cannot sample actions.") + + action_low = np.asarray(action_space.low, dtype=np.float32) + action_high = np.asarray(action_space.high, dtype=np.float32) + action_shape = action_low.shape + + # Track maximum absolute observed values + tracked_keys = ["joint_velocity", "joint_position", "joint_actuator_force"] + max_observed = {k: 0.0 for k in tracked_keys} + + # Include observation at reset + obs0 = _extract_observations(state) + if isinstance(obs0, dict): + for k in tracked_keys: + if k in obs0: + max_observed[k] = max(max_observed[k], float(np.max(np.abs(np.asarray(obs0[k]))))) + + rng = np.random.RandomState(seed) + for i in range(num_steps): + u = rng.uniform(size=action_shape) + action = action_low + (action_high - action_low) * u + + # Provide a numpy RNG to the env step; wrapper will pass it if accepted. + step_out = env.step(state=state, action=action, rng=env.make_rng(seed + i + 1)) + + # Unpack next state from common return conventions + if hasattr(step_out, "state"): + next_state = step_out.state + elif isinstance(step_out, (tuple, list)) and len(step_out) >= 1: + next_state = step_out[0] + else: + next_state = step_out + + obs = _extract_observations(next_state) + if isinstance(obs, dict): + for k in tracked_keys: + if k in obs: + val = float(np.max(np.abs(np.asarray(obs[k])))) + if val > max_observed[k]: + max_observed[k] = val + + state = next_state + + # Print recommended bounds with a 20% safety margin + print("\n--- Recommended Observation Bounds (20% margin) ---") + for k, v in max_observed.items(): + if v == 0.0: + print(f"{k}: observed max 0.0 (increase sampling or inspect env)") + else: + safe = v * 1.2 + print(f"{k}: [-{safe:.6f}, {safe:.6f}] (observed max: {v:.6f})") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--morphology", type=str, default=None, help="Path to morphology YAML (optional)" + ) + parser.add_argument( + "--num-steps", type=int, default=5000, help="Number of random steps to sample" + ) + parser.add_argument("--seed", type=int, default=42, help="RNG seed") + args = parser.parse_args() + + morph_cfg = load_morphology(args.morphology) + arena_cfg = ArenaConfig() + env_cfg = EnvConfig() + + find_empirical_bounds(morph_cfg, arena_cfg, env_cfg, num_steps=args.num_steps, seed=args.seed) + + +if __name__ == "__main__": + main() diff --git a/scripts/train.py b/scripts/train.py new file mode 100644 index 0000000..c367000 --- /dev/null +++ b/scripts/train.py @@ -0,0 +1,58 @@ +import os +import torch +import hydra +from omegaconf import DictConfig, OmegaConf + +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.configs.register_configs import register_configs +from brittle_star_project.trainers.PPOTrainer import PPOTrainer +from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper +from experiment_logger import init_logger, get_logger + + +def make_env(cfg: BrittleStarConfig) -> BrittleStarJaxEnvWrapper: + """Create the environment using the structured configuration.""" + return BrittleStarJaxEnvWrapper( + morphology=cfg.morphology, + arena=cfg.arena, + env_config=cfg.environment, + num_envs=cfg.ppo.num_envs, + ) + + +@hydra.main(config_path="../configs", config_name="main_config", version_base="1.3") +def main(dict_cfg: DictConfig): + # 1. Convert DictConfig to structured dataclass, ensuring the root schema is applied correctly. + config: BrittleStarConfig = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), dict_cfg) + ) + + # 2. Setup run metadata + # Hydra changes CWD to the output directory by default. + run_dir = os.getcwd() + run_name = os.path.basename(run_dir) + + # 3. Initialize Logger + cfg_dict = OmegaConf.to_container(dict_cfg, resolve=True, throw_on_missing=True) + init_logger( + run_name=run_name, + full_config=cfg_dict, + logging_cfg=config.logging, + base_dir=os.path.dirname(run_dir), + ) + logger = get_logger() + logger.info(f"Hydra-initialized run: {run_name}") + logger.info(f"Output directory: {run_dir}") + + # 4. Setup Environment and Torch + env = make_env(config) + torch.backends.cudnn.deterministic = config.experiment.torch_deterministic + + # 5. Train - pass structured config directly + ppo_trainer = PPOTrainer(config, env, run_dir, run_name) + ppo_trainer.train() + + +if __name__ == "__main__": + register_configs() + main() diff --git a/src/brittle_star_project/MLPs/__init__.py b/src/brittle_star_project/MLPs/__init__.py new file mode 100644 index 0000000..8f9c7f7 --- /dev/null +++ b/src/brittle_star_project/MLPs/__init__.py @@ -0,0 +1,19 @@ +from .mlps import ( + GenericDenseLayersWithActivation, + OneDenseLayerMLP, + Actor, + MessagePasser, + AgentParams, + Storage, +) +from .adjancency_builder import build_adjacency + +__all__ = [ + "GenericDenseLayersWithActivation", + "OneDenseLayerMLP", + "Actor", + "MessagePasser", + "AgentParams", + "Storage", + "build_adjacency", +] diff --git a/src/brittle_star_project/MLPs/adjancency_builder.py b/src/brittle_star_project/MLPs/adjancency_builder.py new file mode 100644 index 0000000..878a02c --- /dev/null +++ b/src/brittle_star_project/MLPs/adjancency_builder.py @@ -0,0 +1,67 @@ +from brittle_star_project.environment.env_config import MorphMode +import jax.numpy as jnp + + +def build_adjacency(segments_per_arm, mode: MorphMode): + num_arms = sum(1 for s in segments_per_arm if s > 0) + num_segments = sum(segments_per_arm) + + # FOR NOW SEMI HARDCODE: + # CENTRALIZED: 1 agent, no stress, adja = 1,1 = [[1]] + # FULLY CONNECTED: 5 agents: adj = alle 1 + # CENTRAL DISK:#arms= 5 agents, only neighbor as adjacent so diagonal kinda.. + # ARM = #segments agents: diago kinda, but extra, center ring too, put center mlps first or.. + + if mode == MorphMode.CENTRALIZED: + return jnp.ones((1, 1)) + + if mode == MorphMode.FULLY_CONNECTED: + adj = jnp.ones((num_arms, num_arms)) # everybody adjacent everybody + return adj + + if mode == MorphMode.RING: # ring + adj = jnp.zeros((num_arms, num_arms)) + for i in range(num_arms): + adj = adj.at[i, i].set(1) # self + adj = adj.at[i, (i - 1) % num_arms].set(1) + adj = adj.at[i, (i + 1) % num_arms].set(1) # left and right.. + return adj + + if mode == MorphMode.SEGMENT: + num_nodes = num_arms + num_segments + adj = jnp.zeros((num_nodes, num_nodes)) + + # first ring + for i in range(num_arms): + # self + adj = adj.at[i, i].set(1) + + # ring neighbors + adj = adj.at[i, (i - 1) % num_arms].set(1) + adj = adj.at[i, (i + 1) % num_arms].set(1) + + # then segment chains + idx = 0 + for arm_idx, seg_count in enumerate(segments_per_arm): + for i in range(seg_count): + seg_node = num_arms + idx + i + + adj = adj.at[seg_node, seg_node].set(1) + if i > 0: + adj = adj.at[seg_node, seg_node - 1].set(1) + if i < seg_count - 1: + adj = adj.at[seg_node, seg_node + 1].set(1) + + idx += seg_count + + idx = 0 + for arm_idx, seg_count in enumerate(segments_per_arm): + first_seg = num_arms + idx # first segment of this arm + + # connect ring node first segment + adj = adj.at[arm_idx, first_seg].set(1) + adj = adj.at[first_seg, arm_idx].set(1) + + idx += seg_count + + return adj diff --git a/src/brittle_star_project/MLPs/mlps.py b/src/brittle_star_project/MLPs/mlps.py new file mode 100644 index 0000000..2c5989f --- /dev/null +++ b/src/brittle_star_project/MLPs/mlps.py @@ -0,0 +1,93 @@ +from dataclasses import dataclass, fields, field + +import flax.linen as nn +import jax.numpy as jnp +import jax.tree_util +from typing import Sequence, Callable +from flax.linen.initializers import constant, orthogonal +from flax.core import FrozenDict + + +# semi generic so we can easily make a config for it in experiments +class GenericDenseLayersWithActivation(nn.Module): + layer_sizes: Sequence[int] = field(default_factory=lambda: [64, 64]) + activation: Callable = nn.tanh + + @nn.compact + def __call__(self, x): + for size in self.layer_sizes: + x = nn.Dense(size, kernel_init=orthogonal(jnp.sqrt(2)))(x) + x = self.activation(x) + return x + + +class OneDenseLayerMLP(nn.Module): + @nn.compact + def __call__(self, x): + return nn.Dense(1, kernel_init=orthogonal(1), bias_init=constant(0.0))(x) + + +class Actor(nn.Module): + action_dim: int + + @nn.compact + def __call__(self, x): + mean = nn.Dense(self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0))(x) + log_std = self.param("log_std", nn.initializers.zeros, (self.action_dim,)) + return mean, log_std + + +class MessagePasser(nn.Module): + hidden_dim: int + num_propagation_steps: int + adj_matrix: jnp.ndarray + + @nn.compact + def __call__(self, x: jnp.ndarray): + for _ in range(self.num_propagation_steps): + # (n_nodes, feat) + messages = nn.Dense(self.hidden_dim)(x) + messages = nn.tanh(messages) + + # note: if mean is wanted: adj_matrix / (adj.sum(axis=-1, keepdims=True) + 1e-8) + agg = self.adj_matrix + aggregated = agg @ messages + + x_concat = jnp.concatenate([x, aggregated], axis=-1) + + gate = nn.sigmoid(nn.Dense(self.hidden_dim)(x_concat)) + candidate = nn.tanh(nn.Dense(self.hidden_dim)(x_concat)) + x = gate * x + (1 - gate) * candidate + + return x + + +@jax.tree_util.register_dataclass +@dataclass +class AgentParams: + sensor_params: FrozenDict | dict + actor_params: FrozenDict | dict + critic_params: FrozenDict | dict + feature_extractor_params: FrozenDict | dict + message_passer_params: FrozenDict | dict + + +@jax.tree_util.register_dataclass +@dataclass +class Storage: + obs: jnp.ndarray + actions: jnp.ndarray + logprobs: jnp.ndarray + dones: jnp.ndarray + values: jnp.ndarray + advantages: jnp.ndarray + returns: jnp.ndarray + rewards: jnp.ndarray + + raw_actions: jnp.ndarray | None = None # before clipping + means: jnp.ndarray | None = None # policy mean + stds: jnp.ndarray | None = None # policy std + + def replace(self, **kwargs) -> "Storage": + fs = fields(self) + return Storage(**{f.name: kwargs.get(f.name, getattr(self, f.name)) for f in fs}) diff --git a/src/brittle_star_project/MLPs/routing.py b/src/brittle_star_project/MLPs/routing.py new file mode 100644 index 0000000..4033cb9 --- /dev/null +++ b/src/brittle_star_project/MLPs/routing.py @@ -0,0 +1,22 @@ +"""Shared JAX routing utilities for decentralized multi-agent models.""" + +import jax + + +def apply_per_node(apply_fn, params, x): + """Apply a Flax module independently to each node. + + Args: + apply_fn: The module's ``apply`` method (e.g. ``sensor.apply``). + params: Per-node parameters with shape ``(num_nodes, ...)``. + x: Input tensor with shape ``(batch, num_nodes, features)``. + + Returns: + Output tensor with shape ``(batch, num_nodes, out_features)``. + """ + + def apply_single_node(p, x_node): + # x_node: (batch, feat) — one node's input across the batch + return jax.vmap(lambda xi: apply_fn(p, xi))(x_node) + + return jax.vmap(apply_single_node, in_axes=(0, 1), out_axes=1)(params, x) diff --git a/src/brittle_star_project/__init__.py b/src/brittle_star_project/__init__.py new file mode 100644 index 0000000..4eec766 --- /dev/null +++ b/src/brittle_star_project/__init__.py @@ -0,0 +1,28 @@ +from .environment.env_types import Backend, Task +from .environment.env_config import ArenaConfig, EnvConfig, MorphologyConfig +from .environment.factory import BrittleStarEnvFactory +from .environment.env_wrapper import BrittleStarEnv +from .evaluation import ( + PolicyAgent, + ControlPolicy, + load_metadata, + rollout_headless, + rollout_viewer, + EpisodeResult, +) + +__all__ = [ + "ArenaConfig", + "Backend", + "BrittleStarEnv", + "BrittleStarEnvFactory", + "EnvConfig", + "MorphologyConfig", + "Task", + "PolicyAgent", + "ControlPolicy", + "load_metadata", + "rollout_headless", + "rollout_viewer", + "EpisodeResult", +] diff --git a/src/brittle_star_project/configs/config_architecture.py b/src/brittle_star_project/configs/config_architecture.py new file mode 100644 index 0000000..e8cf5a5 --- /dev/null +++ b/src/brittle_star_project/configs/config_architecture.py @@ -0,0 +1,66 @@ +from dataclasses import dataclass, field +from typing import List, Optional + + +@dataclass +class LayerConfig: + hidden_dims: List[int] = field(default_factory=lambda: [64, 64]) + activation: str = "tanh" + + +@dataclass +class ArchitectureConfig: + """Base class for actor-critic network configurations. + + Both centralized and decentralized architectures share a centralized critic + composed of a feature extractor followed by a shallow output layer. + + See docs/design/actor-critic.md for the full design rationale. + """ + + name: str = "base" + + # Actor pipeline + sensor: Optional[LayerConfig] = None + propagator: Optional[LayerConfig] = None + motor: Optional[LayerConfig] = None + + # Critic pipeline + feature_extractor: Optional[LayerConfig] = None + critic: Optional[LayerConfig] = None + + # Decentralized + message_passing_steps: Optional[int] = None + topology_type: Optional[str] = None # Supported values: "ring", "fully_connected" + + +@dataclass +class CentralizedConfig(ArchitectureConfig): + """Centralized actor-critic architecture (baseline). + + The actor is a single global policy composed of a sensor (input network) + and a motor (output network). The sensor receives the full concatenated + global observation; the motor projects the hidden state to all joint actions. + + See docs/design/actor-critic.md for the full design rationale. + """ + + name: str = "centralized" + + +@dataclass +class DecentralizedConfig(ArchitectureConfig): + """Decentralized actor architecture (NerveNet-MLP variant). + + Each node runs a local sensor, exchanges messages with neighbours via a + propagator for a fixed number of steps, and then a local motor produces + the joint offset for that node only. + + The critic remains centralized (shared with the base class): it receives the + full concatenated global observation and outputs a single scalar. + + See docs/design/actor-critic.md and docs/design/communication.md for the + full design rationale. + """ + + name: str = "decentralized" diff --git a/src/brittle_star_project/configs/config_evaluation.py b/src/brittle_star_project/configs/config_evaluation.py new file mode 100644 index 0000000..c9c248b --- /dev/null +++ b/src/brittle_star_project/configs/config_evaluation.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass +class EvaluationConfig: + """Evaluation settings. + + Currently used for synchronous checkpoint evaluation during training. + """ + + # When enabled, each saved checkpoint is evaluated headlessly and the results + # are appended to a CSV in the run's metrics/ folder. + evaluate_checkpoints: bool = False + eval_max_steps: int = 5000 + eval_seed: int = 0 + + # Cross-model comparison settings. + # comparison_base_seed is the starting seed for generating episode seeds. + comparison_base_seed: int = 0 + # comparison_num_episodes controls how many target positions to evaluate for each model. + comparison_num_episodes: int = 5 + # comparison_models lists the paths (relative to workspace root) to the .cleanrl_model files. + comparison_models: list[str] = field(default_factory=list) + # Path where the comparison results CSV will be saved (relative to workspace root). + comparison_output_csv: str = "metrics/model_comparison.csv" + # Morphology override YAML paths for cross-morphology comparison. + # Each path points to a file in configs/morphology/ (e.g., "configs/morphology/3_arms.yaml"). + # When empty, each model is evaluated only on its training morphology. + comparison_morphologies: list[str] = field(default_factory=list) + + def __post_init__(self) -> None: + if self.evaluate_checkpoints and self.eval_max_steps <= 0: + raise ValueError( + "Configuration Error: 'eval_max_steps' must be > 0 when " + "'evaluate_checkpoints' is enabled." + ) diff --git a/src/brittle_star_project/configs/config_experiment.py b/src/brittle_star_project/configs/config_experiment.py new file mode 100644 index 0000000..7409811 --- /dev/null +++ b/src/brittle_star_project/configs/config_experiment.py @@ -0,0 +1,11 @@ +from dataclasses import dataclass + + +@dataclass +class ExperimentConfig: + exp_name: str = "brittle_star_ppo" + seed: int = 1 + torch_deterministic: bool = True + cuda: bool = True + debug_sanity: bool = False + base_run_dir: str = "runs" diff --git a/src/brittle_star_project/configs/config_ppo.py b/src/brittle_star_project/configs/config_ppo.py new file mode 100644 index 0000000..d0a29bf --- /dev/null +++ b/src/brittle_star_project/configs/config_ppo.py @@ -0,0 +1,22 @@ +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class PPOConfig: + learning_rate: float = 2.5e-4 + total_timesteps: int = 10000000 + num_envs: int = 100 + num_steps: int = 128 + anneal_lr: bool = True + gamma: float = 0.99 + gae_lambda: float = 0.95 + num_minibatches: int = 4 + update_epochs: int = 4 + norm_adv: bool = True + clip_coef: float = 0.1 + clip_vloss: bool = True + ent_coef: float = 0.01 + vf_coef: float = 0.5 + max_grad_norm: float = 0.5 + target_kl: Optional[float] = None diff --git a/src/brittle_star_project/configs/config_simulation.py b/src/brittle_star_project/configs/config_simulation.py new file mode 100644 index 0000000..0a75b8d --- /dev/null +++ b/src/brittle_star_project/configs/config_simulation.py @@ -0,0 +1,35 @@ +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class SimulationSettings: + """Settings for the simulation script.""" + + model_path: Optional[str] = None + + # Script behavior + headless: bool = False + # If None, viewer mode runs until window closed or target reached. + max_steps: Optional[int] = None + + # Override morphology for amputation experiments. + # When set, the environment uses this morphology instead of the trained one. + # Points to a morphology config YAML file (e.g. configs/morphology/3_arms.yaml). + # Observations are padded from the override morphology UP TO the training + # morphology's shape via compute_padding_masks(override, reference=training). + morphology_override: Optional[str] = None + + # Video recording (requires [evaluation] extra) + record_video: bool = False + # When None, video is saved in a per-model evaluation folder alongside the model. + video_output_path: Optional[str] = None + # Camera ID to use for video recording (1 is usually the close-up camera) + camera_id: int = 1 + video_width: int = 640 + video_height: int = 480 + video_fps: int = 60 + + # Optional override for the sidecar metadata YAML file. + # If None, it defaults to the model_path with a `_metadata.yaml` suffix. + metadata_path: Optional[str] = None diff --git a/src/brittle_star_project/configs/main_config.py b/src/brittle_star_project/configs/main_config.py new file mode 100644 index 0000000..7caa1aa --- /dev/null +++ b/src/brittle_star_project/configs/main_config.py @@ -0,0 +1,36 @@ +from dataclasses import dataclass, field + +from experiment_logger.config_logger import LoggingConfig +from brittle_star_project.configs.config_experiment import ExperimentConfig +from brittle_star_project.configs.config_evaluation import EvaluationConfig +from brittle_star_project.configs.config_ppo import PPOConfig +from brittle_star_project.configs.config_architecture import ArchitectureConfig +from brittle_star_project.configs.config_simulation import SimulationSettings +from brittle_star_project.environment.env_config import ( + MorphologyConfig, + ArenaConfig, + EnvConfig, + ObservationBoundsConfig, +) + + +@dataclass +class BrittleStarConfig: + """Root configuration for a brittle star training run. + + Composed of strictly separated sub-configs. Each sub-config can be swapped + independently via CLI or a different YAML file. See configs/README.md. + """ + + experiment: ExperimentConfig = field(default_factory=ExperimentConfig) + logging: LoggingConfig = field(default_factory=LoggingConfig) + evaluation: EvaluationConfig = field(default_factory=EvaluationConfig) + ppo: PPOConfig = field(default_factory=PPOConfig) + # This field is polymorphic; defaults to the base class to allow subclasses + # (CentralizedConfig, DecentralizedConfig) to be merged in via Hydra. + architecture: ArchitectureConfig = field(default_factory=ArchitectureConfig) + morphology: MorphologyConfig = field(default_factory=MorphologyConfig) + arena: ArenaConfig = field(default_factory=ArenaConfig) + environment: EnvConfig = field(default_factory=EnvConfig) + obs_bounds: ObservationBoundsConfig = field(default_factory=ObservationBoundsConfig) + simulation: SimulationSettings = field(default_factory=SimulationSettings) diff --git a/src/brittle_star_project/configs/register_configs.py b/src/brittle_star_project/configs/register_configs.py new file mode 100644 index 0000000..9522212 --- /dev/null +++ b/src/brittle_star_project/configs/register_configs.py @@ -0,0 +1,48 @@ +from hydra.core.config_store import ConfigStore + +from experiment_logger.config_logger import LoggingConfig +from brittle_star_project.configs.config_experiment import ExperimentConfig +from brittle_star_project.configs.config_evaluation import EvaluationConfig +from brittle_star_project.configs.config_ppo import PPOConfig +from brittle_star_project.configs.config_architecture import ( + CentralizedConfig, + DecentralizedConfig, +) +from brittle_star_project.configs.config_simulation import SimulationSettings +from brittle_star_project.environment.env_config import ( + MorphologyConfig, + ArenaConfig, + EnvConfig, + ObservationBoundsConfig, +) +from brittle_star_project.configs.main_config import BrittleStarConfig + + +def register_configs() -> None: + """Register all dataclasses with Hydra's ConfigStore. + + This must be called before hydra.main() processes the config, ensuring + every structured config is validated against its Python schema. Typos in + YAML keys will raise ConfigAttributeError at startup. + """ + cs = ConfigStore.instance() + + # Root schema + cs.store(name="brittle_star_config", node=BrittleStarConfig) + + # Sub-config groups — each group corresponds to a configs/ subdirectory. + cs.store(group="experiment", name="base_experiment", node=ExperimentConfig) + cs.store(group="logging", name="base_logging", node=LoggingConfig) + cs.store(group="evaluation", name="base_evaluation", node=EvaluationConfig) + cs.store(group="ppo", name="base_ppo", node=PPOConfig) + + # Architecture variants — swap via CLI: architecture=decentralized + cs.store(group="architecture", name="centralized_schema", node=CentralizedConfig) + cs.store(group="architecture", name="decentralized_schema", node=DecentralizedConfig) + + # Environment configs + cs.store(group="morphology", name="base_morphology", node=MorphologyConfig) + cs.store(group="arena", name="base_arena", node=ArenaConfig) + cs.store(group="environment", name="base_environment", node=EnvConfig) + cs.store(group="obs_bounds", name="base_obs_bounds", node=ObservationBoundsConfig) + cs.store(group="simulation", name="base_simulation", node=SimulationSettings) diff --git a/src/brittle_star_project/dataclasses/EpisodeStatistics.py b/src/brittle_star_project/dataclasses/EpisodeStatistics.py new file mode 100644 index 0000000..ff2982e --- /dev/null +++ b/src/brittle_star_project/dataclasses/EpisodeStatistics.py @@ -0,0 +1,10 @@ +import flax.struct +import jax.numpy as jnp + + +@flax.struct.dataclass +class EpisodeStatistics: + episode_returns: jnp.ndarray + episode_lengths: jnp.ndarray + returned_episode_returns: jnp.ndarray + returned_episode_lengths: jnp.ndarray diff --git a/src/brittle_star_project/dataclasses/__init__.py b/src/brittle_star_project/dataclasses/__init__.py new file mode 100644 index 0000000..f2d2ad1 --- /dev/null +++ b/src/brittle_star_project/dataclasses/__init__.py @@ -0,0 +1,6 @@ +from .EpisodeStatistics import EpisodeStatistics + + +__all__ = [ + "EpisodeStatistics", +] diff --git a/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py new file mode 100644 index 0000000..cfb60cd --- /dev/null +++ b/src/brittle_star_project/environment/BrittleStarJaxEnvWrapper.py @@ -0,0 +1,114 @@ +import jax +import jax.numpy as jnp + +from experiment_logger import get_logger +from .env_config import EnvConfig, MorphologyConfig, ArenaConfig +from .env_types import Backend +from .factory import BrittleStarEnvFactory +from .padded_obs_wrapper import compute_padding_masks + + +class BrittleStarJaxEnvWrapper: + def __init__( + self, + morphology: MorphologyConfig, + arena: ArenaConfig, + env_config: EnvConfig, + num_envs: int, + backend: Backend = Backend.MJX, + ): + self._morphology = morphology + self._arena = arena + self._env_config = env_config + self._backend = backend + self._num_envs = num_envs + self._env = BrittleStarEnvFactory.create_environment( + self._backend, self._morphology, self._arena, self._env_config + ) + + # Pre-compute masks for observation padding + self._padding_masks = compute_padding_masks(self._morphology.segments_per_arm) + + self._vectorized_reset = jax.jit(jax.vmap(self._env.reset)) + self._vectorized_step = jax.jit(jax.vmap(self._env.step)) + self._vectorized_action_sample = jax.jit(jax.vmap(self._env.action_space.sample)) + + self._action_rng = None + + self.logger = get_logger() + self.logger.info( + f"Initialized BrittleStarJaxEnvWrapper with {num_envs} envs on {backend.value}" + ) + + @property + def backend(self): + return self._backend + + @property + def raw(self): + return self._env + + @property + def padding_masks(self) -> dict: + """Pre-computed boolean masks for amputated limb padding. + + Pass to create_obs_processor so the processor handles padding + after normalization in the correct pipeline order. + """ + return self._padding_masks + + @property + def single_action_space(self): + return self._env.action_space + + @property + def single_observation_space(self): + return self._env.observation_space + + def reset(self, seed: int = 0, target_position: tuple[float, float] | None = None): + self.logger.info(f"Resetting vectorized environment environments with seed {seed}") + self._action_rng, env_rng = jax.random.split(jax.random.PRNGKey(seed), 2) + env_rngs = jnp.array(jax.random.split(env_rng, self._num_envs)) + + # If a target_position is provided, pass it through to the underlying env.reset + if target_position is None: + state = jax.jit(jax.vmap(lambda rng: self._env.reset(rng=rng)))(env_rngs) + else: + tp = jnp.asarray(target_position, dtype=jnp.float32) + tp_batched = jnp.tile(tp[None, :], (self._num_envs, 1)) + state = jax.jit(jax.vmap(lambda rng, t: self._env.reset(rng=rng, target_position=t)))( + env_rngs, tp_batched + ) + + return state + + def sample_actions(self): + assert self._action_rng is not None, "Call reset() before sample_actions()" + self._action_rng, *sub_rngs = jnp.array( + jax.random.split(self._action_rng, self._num_envs + 1) + ) + return self._vectorized_action_sample(rng=jnp.array(sub_rngs)) + + def step(self, state, action): + return self._vectorized_step(state=state, action=action) + + def close(self): + self._env.close() + + @staticmethod + def default(num_envs: int, backend: Backend = Backend.MJX) -> "BrittleStarJaxEnvWrapper": + morphology = MorphologyConfig() + arena = ArenaConfig() + env_config = EnvConfig() + return BrittleStarJaxEnvWrapper( + morphology, arena, env_config, num_envs=num_envs, backend=backend + ) + + def __str__(self): + morphology_str = str(self._morphology) + arena_str = str(self._arena) + env_config_str = str(self._env_config) + return ( + f"BrittleStarJaxEnvWrapper(backend={self._backend}, num_envs={self._num_envs}, " + + f"morphology={morphology_str}, arena={arena_str}, env_config={env_config_str})" + ) diff --git a/src/brittle_star_project/environment/__init__.py b/src/brittle_star_project/environment/__init__.py new file mode 100644 index 0000000..7a9a9eb --- /dev/null +++ b/src/brittle_star_project/environment/__init__.py @@ -0,0 +1,19 @@ +from .env_config import ArenaConfig, EnvConfig, MorphologyConfig, MorphMode +from .env_types import Backend, Task +from .env_wrapper import BrittleStarEnv +from .factory import BrittleStarEnvFactory +from .obs_processing import create_obs_processor +from .padded_obs_wrapper import compute_padding_masks + +__all__ = [ + "ArenaConfig", + "EnvConfig", + "MorphologyConfig", + "Backend", + "Task", + "BrittleStarEnv", + "BrittleStarEnvFactory", + "MorphMode", + "create_obs_processor", + "compute_padding_masks", +] diff --git a/src/brittle_star_project/environment/env_config.py b/src/brittle_star_project/environment/env_config.py new file mode 100644 index 0000000..33e8bdf --- /dev/null +++ b/src/brittle_star_project/environment/env_config.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum + +from .env_types import Task + + +class MorphMode(Enum): + CENTRALIZED = 0 + FULLY_CONNECTED = 1 + RING = 2 + SEGMENT = 3 + + +@dataclass +class MorphologyConfig: + """Brittle star morphology configuration. + + segments_per_arm defines the number of segments for each arm. The length of + this list implicitly sets the number of arms. Use 0 segments to represent + a fully amputated arm (e.g., [4, 0, 4, 2, 4] for a 5-arm morphology with + arm 1 removed and arm 3 shortened). + + The upstream biorobot library natively supports per-arm segment counts. + """ + + segments_per_arm: list[int] = field(default_factory=lambda: [4, 4, 4, 4, 4]) + use_p_control: bool = True + use_torque_control: bool = False + morph_mode: MorphMode = MorphMode.CENTRALIZED + + @property + def num_arms(self) -> int: + return len(self.segments_per_arm) + + +@dataclass +class ArenaConfig: + size: list[float] = field(default_factory=lambda: [10.0, 5.0]) + sand_ground_color: bool = True + attach_target: bool = True + wall_height: float = 1.5 + wall_thickness: float = 0.1 + + +@dataclass +class EnvConfig: + """Shared environment settings. + + Note: Some tasks have additional parameters (see fields below). + """ + + task: Task = Task.DIRECTED_LOCOMOTION + + simulation_time: float = 10000.0 + num_physics_steps_per_control_step: int = 10 + time_scale: int = 2 + + camera_ids: list[int] = field(default_factory=lambda: [0, 1]) + # (height, width) + render_size: list[int] = field(default_factory=lambda: [480, 640]) + + joint_randomization_noise_scale: float = 0.0 + + # Directed locomotion + target_distance: float = 3.0 + + # Light escape + # Per docs in upstream env config: integer factors of 200. + light_perlin_noise_scale: int = 0 + + +@dataclass +class ObservationBoundsConfig: + """Physical observation bounds for deterministic min-max normalization.""" + + # Empirical testing based on the extract_observation_bounds.py script run for 1.000.000 steps + + # Based on max. ctrlrange (0.78539816339744828) in XML, but empirical testing went slightly over + joint_position: list[float] = field(default_factory=lambda: [-0.8, 0.8]) + # Empirical testing showed max. 3.22, adding buffer to be safe. Consider higher values "fast". + joint_velocity: list[float] = field(default_factory=lambda: [-5.0, 5.0]) + # Based on max. forceRange in XML, verified with empirical testing + joint_actuator_force: list[float] = field(default_factory=lambda: [-3.75, 3.75]) + # Based on intuition and reasoning + segment_contact: list[float] = field(default_factory=lambda: [0.0, 1.0]) + robot_direction_to_target: list[float] = field(default_factory=lambda: [-1.0, 1.0]) + disk_z_tilt: list[float] = field(default_factory=lambda: [0.0, 3.141592653589793]) + + def to_bounds_dict(self) -> dict[str, tuple[float, float]]: + return { + "disk_z_tilt": tuple(self.disk_z_tilt), + "joint_actuator_force": tuple(self.joint_actuator_force), + "joint_position": tuple(self.joint_position), + "joint_velocity": tuple(self.joint_velocity), + "robot_direction_to_target": tuple(self.robot_direction_to_target), + "segment_contact": tuple(self.segment_contact), + } diff --git a/src/brittle_star_project/environment/env_types.py b/src/brittle_star_project/environment/env_types.py new file mode 100644 index 0000000..d1be3d5 --- /dev/null +++ b/src/brittle_star_project/environment/env_types.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +from enum import Enum + + +class Backend(str, Enum): + """Physics backend. + + - MJC: MuJoCo C engine + - MJX: MuJoCo XLA (JAX) engine + """ + + MJC = "MJC" + MJX = "MJX" + + +class Task(str, Enum): + """Which brittle-star task/environment to instantiate.""" + + DIRECTED_LOCOMOTION = "directed_locomotion" + LIGHT_ESCAPE = "light_escape" diff --git a/src/brittle_star_project/environment/env_wrapper.py b/src/brittle_star_project/environment/env_wrapper.py new file mode 100644 index 0000000..0d2463c --- /dev/null +++ b/src/brittle_star_project/environment/env_wrapper.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +import inspect +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from .env_config import EnvConfig, MorphologyConfig +from .env_types import Backend + + +@dataclass(slots=True) +class StepResult: + state: Any + reward: float | None = None + terminated: bool | None = None + truncated: bool | None = None + info: dict[str, Any] | None = None + + +class BrittleStarEnv: + """Thin wrapper around the underlying DualMuJoCoEnvironment. + + Goal: hide backend-specific RNG setup and provide a stable place to plug in RL. + """ + + def __init__( + self, + env: Any, + *, + backend: Backend, + config: EnvConfig, + morphology_config: MorphologyConfig | None = None, + ) -> None: + self._env = env + self._backend = backend + self._config = config + self._morphology_config = morphology_config + + @property + def raw(self) -> Any: + return self._env + + @property + def backend(self) -> Backend: + return self._backend + + @property + def config(self) -> EnvConfig: + return self._config + + @property + def morphology_config(self) -> MorphologyConfig | None: + return self._morphology_config + + def make_rng(self, seed: int): + if self._backend == Backend.MJC: + return np.random.RandomState(seed) + + import jax + + return jax.random.PRNGKey(seed) + + def reset(self, *, seed: int = 0, target_position: tuple[float, float, float] | None = None): + rng = self.make_rng(seed) + if target_position is not None: + state = self._env.reset(rng=rng, target_position=target_position) + else: + state = self._env.reset(rng=rng) + return state + + def render(self, *, state: Any): + return self._env.render(state=state) + + def close(self) -> None: + self._env.close() + + def step(self, *, state: Any, action: Any, rng: Any | None = None) -> StepResult: + """Best-effort step wrapper. + + Different env libraries return different tuples; we normalize common cases. + """ + + if not hasattr(self._env, "step"): + raise AttributeError("Underlying env has no step() method") + + step_fn = self._env.step + sig = inspect.signature(step_fn) + params = list(sig.parameters) + + # Common patterns: + # - step(state, action) + # - step(state, action, rng) + # - step(state, action, key) + # We pass rng only if the callable accepts a 3rd arg. + if len(params) >= 3 and rng is not None: + out = step_fn(state, action, rng) + else: + out = step_fn(state, action) + + return out diff --git a/src/brittle_star_project/environment/factory.py b/src/brittle_star_project/environment/factory.py new file mode 100644 index 0000000..523feb2 --- /dev/null +++ b/src/brittle_star_project/environment/factory.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +from dataclasses import asdict + +from moojoco.environment.dual import DualMuJoCoEnvironment + +from .env_config import ArenaConfig, EnvConfig, MorphologyConfig +from .env_types import Backend, Task + + +class BrittleStarEnvFactory: + """Creates brittle-star morphology, arena, and task environment instances.""" + + @staticmethod + def create_morphology(config: MorphologyConfig): + from biorobot.brittle_star.mjcf.morphology.morphology import ( + MJCFBrittleStarMorphology, + ) + from biorobot.brittle_star.mjcf.morphology.specification.default import ( + default_brittle_star_morphology_specification, + ) + + spec = default_brittle_star_morphology_specification( + num_arms=config.num_arms, + num_segments_per_arm=list(config.segments_per_arm), + use_p_control=config.use_p_control, + use_torque_control=config.use_torque_control, + ) + return MJCFBrittleStarMorphology(specification=spec) + + @staticmethod + def create_arena(config: ArenaConfig): + from biorobot.brittle_star.mjcf.arena.aquarium import ( + AquariumArenaConfiguration, + MJCFAquariumArena, + ) + + arena_config = AquariumArenaConfiguration(**asdict(config)) + return MJCFAquariumArena(configuration=arena_config) + + @staticmethod + def create_environment_configuration(config: EnvConfig): + # Import locally so the project can still be imported without these deps. + from biorobot.brittle_star.environment.directed_locomotion.shared import ( + BrittleStarDirectedLocomotionEnvironmentConfiguration, + ) + from biorobot.brittle_star.environment.light_escape.shared import ( + BrittleStarLightEscapeEnvironmentConfiguration, + ) + + common = dict( + joint_randomization_noise_scale=config.joint_randomization_noise_scale, + render_mode="human", + simulation_time=config.simulation_time, + num_physics_steps_per_control_step=config.num_physics_steps_per_control_step, + time_scale=config.time_scale, + camera_ids=config.camera_ids, + render_size=config.render_size, + ) + + match config.task: + case Task.DIRECTED_LOCOMOTION: + return BrittleStarDirectedLocomotionEnvironmentConfiguration( + target_distance=config.target_distance, + **common, + ) + case Task.LIGHT_ESCAPE: + return BrittleStarLightEscapeEnvironmentConfiguration( + light_perlin_noise_scale=config.light_perlin_noise_scale, + **common, + ) + case _: + raise ValueError(f"Unsupported task: {config.task}") + + @staticmethod + def create_environment( + backend: Backend, + morphology_config: MorphologyConfig, + arena_config: ArenaConfig, + env_config: EnvConfig, + ) -> DualMuJoCoEnvironment: + from biorobot.brittle_star.environment.directed_locomotion.dual import ( + BrittleStarDirectedLocomotionEnvironment, + ) + from biorobot.brittle_star.environment.light_escape.dual import ( + BrittleStarLightEscapeEnvironment, + ) + + morphology = BrittleStarEnvFactory.create_morphology(morphology_config) + arena = BrittleStarEnvFactory.create_arena(arena_config) + env_configuration = BrittleStarEnvFactory.create_environment_configuration(env_config) + + match env_config.task: + case Task.DIRECTED_LOCOMOTION: + env_class = BrittleStarDirectedLocomotionEnvironment + case Task.LIGHT_ESCAPE: + env_class = BrittleStarLightEscapeEnvironment + case _: + raise ValueError(f"Unsupported task: {env_config.task}") + + env = env_class.from_morphology_and_arena( + morphology=morphology, + arena=arena, + configuration=env_configuration, + backend=backend.value, + ) + + from experiment_logger import get_logger + + get_logger().info(f"Created {env_config.task.value} env on backend {backend.value}") + + return env diff --git a/src/brittle_star_project/environment/obs_processing.py b/src/brittle_star_project/environment/obs_processing.py new file mode 100644 index 0000000..e169ee3 --- /dev/null +++ b/src/brittle_star_project/environment/obs_processing.py @@ -0,0 +1,192 @@ +import jax +import jax.numpy as jnp +from typing import Dict, Tuple, Optional + +from brittle_star_project.environment.env_config import MorphMode + +from experiment_logger import get_logger + +logger = get_logger() + +_JOINT_SCALED_KEYS = frozenset( + { + "joint_position", + "joint_velocity", + "joint_actuator_force", + "actuator_force", + } +) + +_SEGMENT_SCALED_KEYS = frozenset( + { + "segment_contact", + } +) + + +def _build_joint_indices(segments_per_arm, indices_mlp): + indices = [] + start = 0 + for i, segs in enumerate(segments_per_arm): + # 2 joints per segment + if i in indices_mlp: + count = segs * 2 + idx = jnp.arange(start, start + count) + indices.append(idx) + start += count + return indices + + +def _build_segment_indices(segments_per_arm, indices_mlp): + indices = [] + start = 0 + for i, segs in enumerate(segments_per_arm): + if i in indices_mlp: + idx = jnp.arange(start, start + segs) + indices.append(idx) + start += segs + return indices + + +def create_obs_processor( + bounds_dict: Dict[str, Tuple[float, float]], + num_arms: int, + needed_copies: int, + padding_masks: Optional[Dict] = None, + morph_mode: MorphMode = MorphMode.CENTRALIZED, + segments_per_arm=[4, 4, 4, 4, 4], + agent_indices=[0, 1, 2, 3, 4], +): + # made a set to allow O(1) search + ordered_keys = frozenset( + [ + "disk_z_tilt", + "joint_actuator_force", + "joint_position", + "joint_velocity", + "robot_direction_to_target", + "segment_contact", + ] + ) + segment_indices = _build_segment_indices(segments_per_arm, agent_indices) + joint_indices = _build_joint_indices(segments_per_arm, agent_indices) + + def _add_derived_features(obs: dict) -> dict: + new_obs = dict(obs) + if "disk_rotation" in new_obs: + rot = new_obs["disk_rotation"] + new_obs["disk_z_tilt"] = jnp.sqrt(jnp.pow(rot[0], 2) + jnp.pow(rot[1], 2)) + + if "unit_xy_direction_to_target" in new_obs: + yaw = rot[2] + unit_x, unit_y = new_obs["unit_xy_direction_to_target"] + cos_yaw, sin_yaw = jnp.cos(yaw), jnp.sin(yaw) + new_x = unit_x * cos_yaw + unit_y * sin_yaw + new_y = -unit_x * sin_yaw + unit_y * cos_yaw + new_obs["robot_direction_to_target"] = jnp.stack([new_x, new_y]) + + return new_obs + + def _normalize_features(obs: dict) -> dict: + normalized = {} + for key, arr in obs.items(): + if key in bounds_dict: + low, high = bounds_dict[key] + if low == -1.0 and high == 1.0: + normalized[key] = jnp.clip(arr, -1.0, 1.0) + else: + arr_clipped = jnp.clip(arr, low, high) + normalized[key] = 2.0 * (arr_clipped - low) / (high - low) - 1.0 + else: + normalized[key] = arr + return normalized + + def _split_to_agents(obs: dict, morph_mode) -> dict: + output = {} + num_agents = needed_copies # IMPORTANT: number of MLPs + + segs_per_arm = 4 + joints_per_segment = 2 + joints_per_arm = segs_per_arm * joints_per_segment + for key, arr in obs.items(): + arr = jnp.asarray(arr) + if arr.size == 0: + continue + + if arr.ndim == 0: + arr = arr.reshape(1) + + if key in _SEGMENT_SCALED_KEYS: + per_agent = [] + for i, _ in enumerate(agent_indices): + idx = segment_indices[i] + taken = jnp.take(arr, idx, axis=0) + pad_len = segs_per_arm - taken.shape[0] + padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1)) + + per_agent.append(padded.reshape(-1)) + arr = jnp.stack(per_agent) + elif key in _JOINT_SCALED_KEYS: + per_agent = [] + for i, _ in enumerate(agent_indices): + idx = joint_indices[i] + taken = jnp.take(arr, idx, axis=0) + pad_len = joints_per_arm - taken.shape[0] + padded = jnp.pad(taken, [(0, pad_len)] + [(0, 0)] * (taken.ndim - 1)) + + per_agent.append(padded.reshape(-1)) + arr = jnp.stack(per_agent) + else: + arr = jnp.repeat(arr[None, :], num_agents, axis=0) + + if morph_mode == MorphMode.CENTRALIZED: + output[key] = arr.reshape(1, -1) + elif key in _JOINT_SCALED_KEYS: + output[key] = arr.reshape(num_agents, -1) + elif key in _SEGMENT_SCALED_KEYS: + output[key] = arr[:, None] + else: + output[key] = arr + + return output + + def _flatten_features(obs: dict) -> jnp.ndarray: + """ + Input: + key -> (num_arms, feat_per_key) + + Output: + (num_arms, total_features) + """ + values = [] + + for key in sorted(ordered_keys): + if key not in obs: + continue + + arr = jnp.asarray(obs[key]) # (num_arms, feat) + + if arr.size == 0: + continue + + if arr.ndim == 1: + arr = arr[:, None] + + arr = arr.reshape(arr.shape[0], -1) + + values.append(arr) + + return jnp.concatenate(values, axis=-1) # (num_arms, total_feat) + + def _process_single(obs_dict: dict) -> jnp.ndarray: + processed = _add_derived_features(obs_dict) + processed = _normalize_features(processed) + processed = _split_to_agents(processed, morph_mode) + flat = _flatten_features(processed) # (num_arms, total_feat) + + logger.debug(f"[FLATTENED FINAL] shape: {flat.shape}") + logger.debug(f"[PER AGENT] example row 0 shape: {flat[0].shape}") + + return flat # (agents, feat) + + return jax.jit(jax.vmap(_process_single)) diff --git a/src/brittle_star_project/environment/padded_obs_wrapper.py b/src/brittle_star_project/environment/padded_obs_wrapper.py new file mode 100644 index 0000000..3216dfe --- /dev/null +++ b/src/brittle_star_project/environment/padded_obs_wrapper.py @@ -0,0 +1,54 @@ +"""Observation padding masks for amputated brittle star morphologies.""" + +from __future__ import annotations + +from typing import Any, Sequence + +import jax.numpy as jnp + + +def compute_padding_masks( + segments_per_arm: Sequence[int], + reference_segments_per_arm: Sequence[int] = (4, 4, 4, 4, 4), +) -> dict[str, Any]: + """Pre-compute boolean masks for spatial insertion of observations. + + Args: + segments_per_arm: The current (possibly amputated) morphology. + reference_segments_per_arm: The full morphology that defines the expected size. + + Returns: + A dict containing 1D boolean masks and target sizes. + """ + if len(segments_per_arm) != len(reference_segments_per_arm): + raise ValueError( + f"Morphology mismatch: current has {len(segments_per_arm)} arms, " + f"but reference requires {len(reference_segments_per_arm)} arms." + ) + + mask_1x = [] + mask_2x = [] + + for arm_idx, (actual, ref) in enumerate(zip(segments_per_arm, reference_segments_per_arm)): + if not isinstance(actual, int): + actual = actual.item() + + if not isinstance(ref, int): + ref = ref.item() + + if not (0 <= actual <= ref): + raise ValueError( + f"Invalid amputation at arm {arm_idx}: " + f"actual segments ({actual}) must be between 0 and reference ({ref})." + ) + # 1x scaling (e.g., contacts: 1 value per segment) + mask_1x.extend([True] * actual + [False] * (ref - actual)) + # 2x scaling (e.g., joints: 2 values per segment) + mask_2x.extend([True] * (actual * 2) + [False] * ((ref - actual) * 2)) + + return { + "mask_1x": jnp.array(mask_1x, dtype=bool), + "mask_2x": jnp.array(mask_2x, dtype=bool), + "target_size_1x": sum(reference_segments_per_arm), + "target_size_2x": sum(reference_segments_per_arm) * 2, + } diff --git a/src/brittle_star_project/evaluation/__init__.py b/src/brittle_star_project/evaluation/__init__.py new file mode 100644 index 0000000..e529115 --- /dev/null +++ b/src/brittle_star_project/evaluation/__init__.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from .checkpoint import load_metadata, load_params, metadata_to_configs, TrainingConfig +from .evaluate_mjx import ( + CheckpointEvalResult, + append_checkpoint_eval_row, + build_eval_rollout_fn, + evaluate_checkpoint_mjx, +) +from .evaluate import evaluate_policy +from .policy import PolicyAgent, ControlPolicy +from .rollout import rollout_headless, rollout_viewer, EpisodeResult +from .video import record_episode, create_evaluation_dir, save_evaluation_metadata +from .eval_env_builder import EvalEnvBundle, build_eval_env + +__all__ = [ + # checkpoint loading + "load_metadata", + "load_params", + "metadata_to_configs", + "TrainingConfig", + # MJX evaluation + "CheckpointEvalResult", + "append_checkpoint_eval_row", + "build_eval_rollout_fn", + "evaluate_checkpoint_mjx", + # CPU evaluation + "evaluate_policy", + # policy + "PolicyAgent", + "ControlPolicy", + # rollout + "rollout_headless", + "rollout_viewer", + "EpisodeResult", + # video + "record_episode", + "create_evaluation_dir", + "save_evaluation_metadata", + # env builder + "EvalEnvBundle", + "build_eval_env", +] diff --git a/src/brittle_star_project/evaluation/checkpoint.py b/src/brittle_star_project/evaluation/checkpoint.py new file mode 100644 index 0000000..9b868d5 --- /dev/null +++ b/src/brittle_star_project/evaluation/checkpoint.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +import yaml +from dataclasses import dataclass +from pathlib import Path + +from collections.abc import Mapping + +import flax +from omegaconf import OmegaConf + +from brittle_star_project.environment.env_config import ( + MorphologyConfig, + ArenaConfig, + EnvConfig, + ObservationBoundsConfig, +) + + +@dataclass +class TrainingConfig: + """Holds typed configurations extracted from a training run's metadata.""" + + morphology: MorphologyConfig + arena: ArenaConfig + environment: EnvConfig + obs_bounds: ObservationBoundsConfig + + +def load_params(path: Path) -> dict: + """Load model parameters from a .flax checkpoint file.""" + payload = path.read_bytes() + restored = flax.serialization.msgpack_restore(payload) + + sensor_params = None + actor_params = None + message_passer_params = None + + # Extract params from restored checkpoint + if isinstance(restored, Mapping): + params_sub = restored.get("params", {}) + sensor_params = restored.get("sensor_params") or params_sub.get("sensor_params") + actor_params = restored.get("actor_params") or params_sub.get("actor_params") + message_passer_params = restored.get("message_passer_params") or params_sub.get( + "message_passer_params" + ) + elif isinstance(restored, (list, tuple)) and len(restored) >= 2: + params_part = restored[1] + if isinstance(params_part, Mapping): + sensor_params = params_part.get("0", params_part.get(0)) + actor_params = params_part.get("1", params_part.get(1)) + elif isinstance(params_part, (list, tuple)) and len(params_part) >= 2: + sensor_params = params_part[0] + actor_params = params_part[1] + + if sensor_params is None or actor_params is None: + raise ValueError(f"Could not extract sensor and actor params from checkpoint: {path}") + + return { + "sensor_params": sensor_params, + "actor_params": actor_params, + "message_passer_params": message_passer_params, + } + + +def load_metadata(model_path: Path, metadata_override_path: Path | None = None) -> dict: + """Discover and load the sidecar metadata YAML file.""" + if metadata_override_path is not None: + metadata_path = metadata_override_path + else: + metadata_path = model_path.with_name(model_path.stem + "_metadata.yaml") + + if not metadata_path.exists(): + raise FileNotFoundError(f"Could not find metadata YAML at {metadata_path}") + with open(metadata_path, "r") as f: + return yaml.safe_load(f) + + +def metadata_to_configs(metadata: dict) -> TrainingConfig: + """Reconstruct typed configuration objects from a metadata dictionary.""" + trained_morphology = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(MorphologyConfig), metadata.get("morphology", {})) + ) + trained_arena = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(ArenaConfig), metadata.get("arena", {})) + ) + + env_dict = metadata.get("environment", {}) + if isinstance(env_dict.get("task"), str): + from brittle_star_project.environment.env_types import Task + + try: + env_dict["task"] = Task[env_dict["task"]].name + except Exception: + try: + env_dict["task"] = Task(env_dict["task"]).name + except Exception: + pass + + trained_environment = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(EnvConfig), env_dict) + ) + trained_obs_bounds = OmegaConf.to_object( + OmegaConf.merge( + OmegaConf.structured(ObservationBoundsConfig), metadata.get("obs_bounds", {}) + ) + ) + + return TrainingConfig( + morphology=trained_morphology, + arena=trained_arena, + environment=trained_environment, + obs_bounds=trained_obs_bounds, + ) diff --git a/src/brittle_star_project/evaluation/eval_env_builder.py b/src/brittle_star_project/evaluation/eval_env_builder.py new file mode 100644 index 0000000..6d45966 --- /dev/null +++ b/src/brittle_star_project/evaluation/eval_env_builder.py @@ -0,0 +1,176 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import jax.numpy as jnp +import numpy as np +import yaml +from omegaconf import OmegaConf + +from brittle_star_project import Backend, BrittleStarEnv, BrittleStarEnvFactory +from brittle_star_project.environment.env_config import MorphMode, MorphologyConfig +from brittle_star_project.environment.obs_processing import create_obs_processor +from brittle_star_project.environment.padded_obs_wrapper import compute_padding_masks +from brittle_star_project.evaluation.checkpoint import TrainingConfig +from brittle_star_project.evaluation.policy import PolicyAgent +from brittle_star_project.MLPs.adjancency_builder import build_adjacency + + +@dataclass +class EvalEnvBundle: + """Everything needed to run a headless evaluation episode.""" + + env: BrittleStarEnv + policy: PolicyAgent + action_low: np.ndarray | None + action_high: np.ndarray | None + action_mask: np.ndarray | None + segments_per_arm: list[int] + num_active_arms: int + architecture: str + + +def build_eval_env( + *, + model_path: Path, + training: TrainingConfig, + metadata: dict, + morphology_override_path: Path | str | None = None, +) -> EvalEnvBundle: + """Build environment + policy for evaluation, optionally with a morphology override.""" + + # 1. Determine environment morphology + if morphology_override_path is not None: + override_path = Path(morphology_override_path) + if not override_path.exists(): + raise FileNotFoundError(f"Could not find morphology override YAML at {override_path}") + with open(override_path, "r") as f: + override_dict = yaml.safe_load(f) + env_morphology = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(MorphologyConfig), override_dict) + ) + # Force morph_mode to be inherited from training since it's baked into weights + env_morphology.morph_mode = training.morphology.morph_mode + else: + env_morphology = training.morphology + + # 2. Build obs_processor with TRAINING morphology padding masks always + padding_masks = compute_padding_masks( + segments_per_arm=env_morphology.segments_per_arm, + reference_segments_per_arm=training.morphology.segments_per_arm, + ) + + training_segs_per_arm = jnp.array(training.morphology.segments_per_arm) + + needed_copies = 0 + agent_indices = [0, 1, 2, 3, 4] + match training.morphology.morph_mode: + case MorphMode.CENTRALIZED: + needed_copies = 1 + case MorphMode.FULLY_CONNECTED | MorphMode.RING: + agent_mask = training_segs_per_arm > 0 + agent_indices = jnp.where(agent_mask)[0].tolist() + needed_copies = jnp.where(training_segs_per_arm > 0, 1, 0).sum().item() + case MorphMode.SEGMENT: + agent_mask = training_segs_per_arm > 0 + agent_indices = jnp.where(agent_mask)[0].tolist() + needed_copies = ( + training_segs_per_arm.sum() + jnp.where(training_segs_per_arm > 0, 1, 0).sum() + ).item() + + num_arms_training = jnp.where(training_segs_per_arm > 0, 1, 0).sum().item() + + obs_processor = create_obs_processor( + bounds_dict=training.obs_bounds.to_bounds_dict(), + padding_masks=padding_masks, + needed_copies=needed_copies, + num_arms=num_arms_training, + morph_mode=training.morphology.morph_mode, + segments_per_arm=env_morphology.segments_per_arm, + agent_indices=agent_indices, + ) + + # 3. Build environment + backend = Backend.MJC + factory = BrittleStarEnvFactory() + raw_env = factory.create_environment( + backend, + env_morphology, + training.arena, + training.environment, + ) + env = BrittleStarEnv( + raw_env, + backend=backend, + config=training.environment, + morphology_config=env_morphology, + ) + + # Calculate the action dimension the model was trained with + training_total_actions = sum(training.morphology.segments_per_arm) * 2 + trained_action_dim = training_total_actions // needed_copies + + # 4. Load policy + message_passing_steps = (metadata.get("architecture", {}) or {}).get("message_passing_steps") + if message_passing_steps is None: + message_passing_steps = 4 + message_passing_steps = int(message_passing_steps) + + adj_matrix = None + if training.morphology.morph_mode != MorphMode.CENTRALIZED: + adj_matrix = build_adjacency( + training.morphology.segments_per_arm, training.morphology.morph_mode + ) + + override_segs = env_morphology.segments_per_arm + if training.morphology.morph_mode in (MorphMode.FULLY_CONNECTED, MorphMode.RING): + for i, segs in enumerate(override_segs): + if segs == 0 and i < adj_matrix.shape[0]: + adj_matrix = adj_matrix.at[i, :].set(0) + adj_matrix = adj_matrix.at[:, i].set(0) + elif training.morphology.morph_mode == MorphMode.SEGMENT: + for i, segs in enumerate(override_segs): + if segs == 0 and i < num_arms_training: + adj_matrix = adj_matrix.at[i, :].set(0) + adj_matrix = adj_matrix.at[:, i].set(0) + + idx = 0 + for arm_idx, seg_count in enumerate(training.morphology.segments_per_arm): + if override_segs[arm_idx] == 0: + for i in range(seg_count): + seg_node = num_arms_training + idx + i + if seg_node < adj_matrix.shape[0]: + adj_matrix = adj_matrix.at[seg_node, :].set(0) + adj_matrix = adj_matrix.at[:, seg_node].set(0) + idx += seg_count + + policy = PolicyAgent.from_checkpoint( + model_path, + action_dim=trained_action_dim, + obs_processor=obs_processor, + message_passing_steps=message_passing_steps, + adj_matrix=adj_matrix, + ) + + # 5. Build action clipping and masks + action_mask = np.asarray(padding_masks["mask_2x"]) + + action_space = getattr(raw_env, "action_space", None) + action_low = ( + None if action_space is None else np.asarray(action_space.low, dtype=np.float32).ravel() + ) + action_high = ( + None if action_space is None else np.asarray(action_space.high, dtype=np.float32).ravel() + ) + + return EvalEnvBundle( + env=env, + policy=policy, + action_low=action_low, + action_high=action_high, + action_mask=action_mask, + segments_per_arm=env_morphology.segments_per_arm, + num_active_arms=sum(1 for s in env_morphology.segments_per_arm if s > 0), + architecture=env_morphology.morph_mode.name, + ) diff --git a/src/brittle_star_project/evaluation/evaluate.py b/src/brittle_star_project/evaluation/evaluate.py new file mode 100644 index 0000000..9d62880 --- /dev/null +++ b/src/brittle_star_project/evaluation/evaluate.py @@ -0,0 +1,58 @@ +"""MJC-based (CPU) checkpoint evaluation. + +This module provides the CPU-bound evaluation path using the standard MJC backend. +It is primarily used by the `evaluate_checkpoints` CLI to compute metrics and +render videos. +""" + +from pathlib import Path + +import numpy as np + +from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper +from brittle_star_project.environment.obs_processing import create_obs_processor +from brittle_star_project.evaluation.policy import PolicyAgent +from brittle_star_project.evaluation.rollout import EpisodeResult, rollout_headless + + +def evaluate_policy( + env: BrittleStarJaxEnvWrapper, + policy_path: str | Path, + seed: int, + max_steps: int, +) -> EpisodeResult: + """Evaluate a trained policy in a CPU-bound environment. + + Args: + env: Initialised CPU environment (MJC backend). + policy_path: Path to the `.cleanrl_model` weights file. + seed: Random seed for environment reset. + max_steps: Maximum number of control steps. + + Returns: + Structured result containing return, length, and distance metrics. + """ + obs_processor = create_obs_processor( + bounds_dict=env.cfg.obs_bounds.to_bounds_dict(), + padding_masks=env.padding_masks, + ) + + action_dim = env.single_action_space.shape[0] + + policy = PolicyAgent.from_checkpoint( + model_path=Path(policy_path), + action_dim=action_dim, + obs_processor=obs_processor, + ) + + action_low = np.asarray(env.single_action_space.low, dtype=np.float32) + action_high = np.asarray(env.single_action_space.high, dtype=np.float32) + + return rollout_headless( + env=env, + policy=policy, + seed=seed, + max_steps=max_steps, + action_low=action_low, + action_high=action_high, + ) diff --git a/src/brittle_star_project/evaluation/evaluate_mjx.py b/src/brittle_star_project/evaluation/evaluate_mjx.py new file mode 100644 index 0000000..f557d2d --- /dev/null +++ b/src/brittle_star_project/evaluation/evaluate_mjx.py @@ -0,0 +1,258 @@ +"""MJX-based headless checkpoint evaluation. + +This module provides a fast, JIT-compiled evaluation path using the MJX +(JAX-accelerated MuJoCo) backend. It is intended for evaluating checkpoints +*during* or *after* a training run, where the environment and policy are +already fully initialised. + +The key functions are: + +- `build_eval_rollout_fn` — builds and JIT-compiles a single-episode rollout function from the + training environment and policy components. +- `evaluate_checkpoint_mjx` — runs that function for a given set of parameters and returns a typed + `CheckpointEvalResult`. +- `append_checkpoint_eval_row` — persists the result to the run's + `metrics/checkpoint_evaluation.csv`, migrating old schemas automatically. +""" + +from __future__ import annotations + +import csv +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable + +import jax +import jax.numpy as jnp + + +@dataclass +class CheckpointEvalResult: + """Structured result from a single MJX checkpoint evaluation episode.""" + + steps: int + """Number of control steps taken (≤ max_steps).""" + + reached_target: bool + """Whether the robot reached the target (terminated) before max_steps.""" + + eval_return: float + """Accumulated shaped reward over the episode.""" + + final_xy_dist: float + """XY distance to target at episode end. 0.0 when ``reached_target`` is True.""" + + initial_xy_dist: float + """XY distance to target at episode start.""" + + +def build_eval_rollout_fn( + *, + env: Any, + obs_processor: Callable, + sensor_apply: Callable, + actor_apply: Callable, + message_passer_apply: Callable | None = None, + action_low: jnp.ndarray, + action_high: jnp.ndarray, + reward_fn: Callable, +) -> Callable: + """Build and JIT-compile a single-episode MJX evaluation rollout. + + All outputs are JAX arrays. Convert to Python scalars before logging. + + Args: + env: The training environment wrapper. Must expose `env.raw` with + `reset` and `step` methods compatible with `jax.vmap`. + obs_processor: Observation normalisation / padding callable, as + returned by `create_obs_processor`. + sensor_apply: The sensor network's `apply` method (JIT-compiled). + actor_apply: The actor network's `apply` method (JIT-compiled). + message_passer_apply: Optional message-passing module apply method. + When provided, it is applied between the sensor and actor, using + `params["message_passer_params"]`. + action_low: Per-joint action lower bound (JAX array, shape `(action_dim,)`). + action_high: Per-joint action upper bound (JAX array, shape `(action_dim,)`). + reward_fn: Shaped reward function with signature + `reward_fn(env_state, next_env_state) -> jnp.ndarray`. + Typically, the module-level `reward_fn` from `PPOTrainer`. + + Returns: + A JIT-compiled callable that runs one deterministic evaluation episode. + """ + # vmap over a batch of 1 so the MJX API is satisfied without any + # extra bookkeeping in the caller. + reset_1 = jax.vmap(env.raw.reset) + step_1 = jax.vmap(env.raw.step) + + def _eval_rollout(params: dict, seed: int, max_steps: int): + rng = jax.random.PRNGKey(seed) + rngs = jnp.asarray(jax.random.split(rng, 1)) + state = reset_1(rng=rngs) + + initial_xy_dist = jnp.squeeze(state.observations["xy_distance_to_target"]) + + t0 = jnp.asarray(0, dtype=jnp.int32) + done0 = jnp.squeeze(state.terminated | state.truncated) + return0 = jnp.asarray(0.0, dtype=jnp.float32) + + def cond(carry): + t, _state, done, _return_ = carry + return jnp.logical_and(t < max_steps, jnp.logical_not(done)) + + def body(carry): + t, state, _done, return_ = carry + + obs = obs_processor(state.observations) + hidden = sensor_apply(params["sensor_params"], obs) + if message_passer_apply is not None: + mp_params = params["message_passer_params"] + hidden = jax.vmap(lambda x: message_passer_apply(mp_params, x))(hidden) + mean, _log_std = actor_apply(params["actor_params"], hidden) + + # Deterministic action: use the actor mean, no exploration noise. + flat_mean = mean.reshape(mean.shape[0], -1) + action = jnp.clip(flat_mean, action_low, action_high) + next_state = step_1(state=state, action=action) + + shaped_reward = reward_fn(state, next_state) + return_ = return_ + jnp.squeeze(shaped_reward) + + done_next = jnp.squeeze(next_state.terminated | next_state.truncated) + return (t + 1, next_state, done_next, return_) + + t, final_state, _done, return_ = jax.lax.while_loop(cond, body, (t0, state, done0, return0)) + + reached_target = jnp.squeeze(final_state.terminated) + final_xy_dist_raw = jnp.squeeze(final_state.observations["xy_distance_to_target"]) + # Clamp to 0 when the target was reached so downstream consumers + # don't have to special-case "terminated" themselves. + final_xy_dist = jnp.where(reached_target, 0.0, final_xy_dist_raw) + + return t, reached_target, return_, final_xy_dist, initial_xy_dist + + return jax.jit(_eval_rollout) + + +def evaluate_checkpoint_mjx( + eval_fn: Callable, + params: dict, + *, + seed: int, + max_steps: int, +) -> CheckpointEvalResult: + """Run one deterministic evaluation episode and return typed metrics. + + Args: + eval_fn: A JIT-compiled function as returned by `build_eval_rollout_fn`. + params: Agent parameter dict (e.g. ``agent_state.params``). + seed: Random seed for environment reset (controls target placement). + max_steps: Maximum number of control steps before the episode is cut off. + + Returns: + A `CheckpointEvalResult` with all JAX arrays converted to + plain Python scalars. + """ + steps, reached, eval_return, final_xy_dist, initial_xy_dist = eval_fn(params, seed, max_steps) + return CheckpointEvalResult( + steps=int(steps), + reached_target=bool(reached), + eval_return=float(eval_return), + final_xy_dist=float(final_xy_dist), + initial_xy_dist=float(initial_xy_dist), + ) + + +_FIELDNAMES = [ + "checkpoint", + "trained_timesteps", + "eval_steps", + "eval_return", + "final_xy_dist", + "initial_xy_dist", + "reached_target", +] + + +def _migrate_csv_if_needed(csv_path: Path) -> None: + """Rewrite the CSV with the canonical field names if the schema changed. + + Best-effort: any exception is silently swallowed so that a schema mismatch + never causes a training crash. + """ + try: + with open(csv_path, "r", newline="") as f: + header = next(csv.reader(f), None) + + if header is None or list(header) == _FIELDNAMES: + return # Nothing to migrate. + + migrated_rows: list[dict[str, Any]] = [] + with open(csv_path, "r", newline="") as f: + for row in csv.DictReader(f): + migrated_rows.append( + { + "checkpoint": row.get("checkpoint", row.get("iteration")), + "trained_timesteps": row.get("trained_timesteps"), + "eval_steps": row.get("eval_steps", row.get("steps_to_target")), + "eval_return": row.get("eval_return"), + "final_xy_dist": row.get("final_xy_dist"), + "initial_xy_dist": row.get("initial_xy_dist"), + "reached_target": row.get("reached_target"), + } + ) + + with open(csv_path, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=_FIELDNAMES) + writer.writeheader() + writer.writerows(migrated_rows) + except Exception: + pass # Never crash training on a migration issue. + + +def append_checkpoint_eval_row( + run_dir: str | Path, + *, + iteration: int, + trained_timesteps: int, + result: CheckpointEvalResult, +) -> Path: + """Append one evaluation row to `/metrics/checkpoint_evaluation.csv`. + + Creates the file (including the `metrics/` directory) if it does not yet + exist. Migrates the file to the current schema if the header has changed. + + Args: + run_dir: Root directory of the training run (Hydra's output dir). + iteration: Training iteration number, used as the checkpoint identifier. + trained_timesteps: Total environment steps taken at this checkpoint. + result: Evaluation result as returned by `evaluate_checkpoint_mjx`. + + Returns: + Absolute path to the CSV file (useful for W&B sync). + """ + metrics_dir = Path(run_dir) / "metrics" + metrics_dir.mkdir(parents=True, exist_ok=True) + csv_path = metrics_dir / "checkpoint_evaluation.csv" + + if csv_path.exists(): + _migrate_csv_if_needed(csv_path) + + file_exists = csv_path.exists() + with open(csv_path, "a", newline="") as f: + writer = csv.DictWriter(f, fieldnames=_FIELDNAMES) + if not file_exists: + writer.writeheader() + writer.writerow( + { + "checkpoint": int(iteration), + "trained_timesteps": int(trained_timesteps), + "eval_steps": result.steps, + "eval_return": result.eval_return, + "final_xy_dist": result.final_xy_dist, + "initial_xy_dist": result.initial_xy_dist, + "reached_target": result.reached_target, + } + ) + + return csv_path diff --git a/src/brittle_star_project/evaluation/policy.py b/src/brittle_star_project/evaluation/policy.py new file mode 100644 index 0000000..00c1e06 --- /dev/null +++ b/src/brittle_star_project/evaluation/policy.py @@ -0,0 +1,168 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any, Protocol + +import jax +import jax.numpy as jnp +import numpy as np + +from brittle_star_project.MLPs.routing import apply_per_node +from brittle_star_project.evaluation.checkpoint import load_params + + +class ControlPolicy(Protocol): + """Protocol for any policy that can produce actions from observations.""" + + def act(self, *, observations: dict[str, Any]) -> np.ndarray: ... + + +class PolicyAgent: + """Wraps a trained Flax actor for deterministic inference.""" + + def __init__( + self, + *, + sensor_params: Any, + actor_params: Any, + message_passer_params: Any | None = None, + message_passing_steps: int | None = None, + adj_matrix: Any | None = None, + action_dim: int, + obs_processor: Any, + ) -> None: + from brittle_star_project.MLPs.mlps import ( + Actor, + GenericDenseLayersWithActivation, + MessagePasser, + ) + + # Infer layer sizes from params + try: + dense_params = ( + sensor_params.get("params", {}) + if isinstance(sensor_params, dict) + else sensor_params["params"] + ) + except Exception: + dense_params = sensor_params + + layer_sizes = [] + idx = 0 + while True: + key = f"Dense_{idx}" + if key not in dense_params: + break + + layer_sizes.append(int(np.asarray(dense_params[key]["kernel"]).shape[-1])) + idx += 1 + + if not layer_sizes: + raise ValueError("Could not infer Dense_* layers from sensor params") + + self._sensor = GenericDenseLayersWithActivation(layer_sizes=layer_sizes) + self._actor = Actor(action_dim=action_dim) + + self._message_passer = None + if message_passer_params is not None and not ( + isinstance(message_passer_params, dict) and len(message_passer_params) == 0 + ): + if message_passing_steps is None or adj_matrix is None: + raise ValueError( + "Checkpoint contains message_passer_params but PolicyAgent was not given " + "message_passing_steps and adj_matrix. Pass these when constructing the agent " + "so decentralized evaluation matches training." + ) + + hidden_dim = int(layer_sizes[-1]) + self._message_passer = MessagePasser( + hidden_dim=hidden_dim, + num_propagation_steps=int(message_passing_steps), + adj_matrix=jnp.asarray(adj_matrix), + ) + self._message_passer.apply = jax.jit(self._message_passer.apply) + self._sensor.apply = jax.jit(self._sensor.apply) + self._actor.apply = jax.jit(self._actor.apply) + self._params = { + "sensor_params": sensor_params, + "actor_params": actor_params, + "message_passer_params": message_passer_params, + } + self._obs_processor = obs_processor + + @classmethod + def from_params( + cls, + *, + sensor_params: Any, + actor_params: Any, + message_passer_params: Any | None = None, + message_passing_steps: int | None = None, + adj_matrix: Any | None = None, + action_dim: int, + obs_processor: Any, + ) -> "PolicyAgent": + """Construct a PolicyAgent directly from in-memory parameters.""" + return cls( + sensor_params=sensor_params, + actor_params=actor_params, + message_passer_params=message_passer_params, + message_passing_steps=message_passing_steps, + adj_matrix=adj_matrix, + action_dim=action_dim, + obs_processor=obs_processor, + ) + + def set_params( + self, + *, + sensor_params: Any, + actor_params: Any, + message_passer_params: Any | None = None, + ) -> None: + """Update parameters for evaluation without rebuilding the model.""" + self._params["sensor_params"] = sensor_params + self._params["actor_params"] = actor_params + self._params["message_passer_params"] = message_passer_params + + @classmethod + def from_checkpoint( + cls, + model_path: Path, + *, + action_dim: int, + obs_processor: Any, + message_passing_steps: int | None = None, + adj_matrix: Any | None = None, + ) -> "PolicyAgent": + """Load params from .flax and construct the agent.""" + params = load_params(model_path) + + return cls( + sensor_params=params["sensor_params"], + actor_params=params["actor_params"], + message_passer_params=params.get("message_passer_params"), + message_passing_steps=message_passing_steps, + adj_matrix=adj_matrix, + action_dim=action_dim, + obs_processor=obs_processor, + ) + + def act(self, *, observations: dict[str, Any]) -> np.ndarray: + """Return deterministic action (actor mean, no exploration noise).""" + batched_obs = jax.tree.map(lambda x: jnp.asarray(x)[None, ...], observations) + obs = self._obs_processor(batched_obs) + + hidden = apply_per_node(self._sensor.apply, self._params["sensor_params"], obs) + + if self._message_passer is not None: + mp_params = self._params.get("message_passer_params") + if mp_params is None or (isinstance(mp_params, dict) and len(mp_params) == 0): + raise ValueError( + "PolicyAgent has a message passer but message_passer_params are missing/empty." + ) + hidden = jax.vmap(lambda x: self._message_passer.apply(mp_params, x))(hidden) + + mean, _log_std = apply_per_node(self._actor.apply, self._params["actor_params"], hidden) + + return np.asarray(mean, dtype=np.float32).ravel() diff --git a/src/brittle_star_project/evaluation/rollout.py b/src/brittle_star_project/evaluation/rollout.py new file mode 100644 index 0000000..f292138 --- /dev/null +++ b/src/brittle_star_project/evaluation/rollout.py @@ -0,0 +1,168 @@ +from __future__ import annotations + +import itertools +import time +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from brittle_star_project import BrittleStarEnv +from brittle_star_project.evaluation.policy import ControlPolicy + + +@dataclass +class EpisodeResult: + return_: float + length: int + reached_target: bool + final_xy_dist: float | None + initial_target_distance: float | None + + +def _get_observations(state: Any) -> dict[str, Any] | None: + return getattr(state, "observations", None) + + +def _get_xy_distance_to_target(observations: dict[str, Any]) -> float | None: + return float(np.asarray(observations["xy_distance_to_target"]).reshape(-1)[0]) + + +def _target_reached(*, state: Any) -> bool: + return bool(getattr(state, "terminated", False) or getattr(state, "truncated", False)) + + +def _maybe_clip_action( + action: np.ndarray, + low: np.ndarray | None, + high: np.ndarray | None, +) -> np.ndarray: + if low is None or high is None: + return action + low = np.asarray(low, dtype=np.float32).ravel() + high = np.asarray(high, dtype=np.float32).ravel() + if low.shape != action.shape or high.shape != action.shape: + return action + return np.clip(action, low, high) + + +def rollout_headless( + *, + env: BrittleStarEnv, + policy: ControlPolicy, + seed: int, + max_steps: int, + action_low: np.ndarray | None, + action_high: np.ndarray | None, + action_mask: np.ndarray | None = None, +) -> EpisodeResult: + """Run an episode headlessly and return the result.""" + state = env.reset(seed=seed) + + ep_return = 0.0 + observations = _get_observations(state) + prev_dist = _get_xy_distance_to_target(observations) if observations else None + initial_target_distance = prev_dist + reached_target = _target_reached(state=state) + + steps = 0 + for _ in range(int(max_steps)): + obs_dict = observations or {} + + action = policy.act(observations=obs_dict) + if action_mask is not None: + action = action[action_mask] + action = _maybe_clip_action(action, action_low, action_high) + + state = env.step(state=state, action=action) + steps += 1 + + observations = _get_observations(state) + cur_dist = _get_xy_distance_to_target(observations) if observations else None + if prev_dist is not None and cur_dist is not None: + ep_return += prev_dist - cur_dist + prev_dist = cur_dist + + reached_target = _target_reached(state=state) + if reached_target: + break + + final_dist = _get_xy_distance_to_target(observations) if observations else None + return EpisodeResult( + return_=ep_return, + length=steps, + reached_target=reached_target, + final_xy_dist=final_dist, + initial_target_distance=initial_target_distance, + ) + + +def rollout_viewer( + *, + env: BrittleStarEnv, + policy: ControlPolicy, + seed: int, + state: Any, + control_dt: float, + max_steps: int | None, + action_low: np.ndarray | None, + action_high: np.ndarray | None, + action_mask: np.ndarray | None = None, +) -> None: + """Run an episode using the interactive MuJoCo viewer.""" + import mujoco.viewer + + model = state.mj_model + data = state.mj_data + + episode_return = 0.0 + observations = _get_observations(state) + + prev_dist = _get_xy_distance_to_target(observations) if observations else None + reached_target = _target_reached(state=state) + + steps = 0 + with mujoco.viewer.launch_passive(model, data) as viewer: + step_iter = range(int(max_steps)) if max_steps is not None else itertools.count() + for _ in step_iter: + if not viewer.is_running(): + break + step_start = time.time() + + obs_dict = observations or {} + + action = policy.act(observations=obs_dict) + if action_mask is not None: + action = action[action_mask] + action = _maybe_clip_action(action, action_low, action_high) + + with viewer.lock(): + state = env.step(state=state, action=action) + + if not viewer.is_running(): + break + viewer.sync() + + steps += 1 + + observations = _get_observations(state) + cur_dist = _get_xy_distance_to_target(observations) if observations else None + if prev_dist is not None and cur_dist is not None: + episode_return += prev_dist - cur_dist + prev_dist = cur_dist + + reached_target = _target_reached(state=state) + if reached_target: + break + + remaining = control_dt - (time.time() - step_start) + if remaining > 0: + time.sleep(remaining) + + dist = _get_xy_distance_to_target(observations) if observations else None + dist_str = "n/a" if dist is None else f"{dist:.3f}" + print( + "episode done: " + f"return={episode_return:.6f}, len={steps}, " + f"target_reached={reached_target}, final_xy_dist={dist_str}" + ) diff --git a/src/brittle_star_project/evaluation/video.py b/src/brittle_star_project/evaluation/video.py new file mode 100644 index 0000000..cf2e5c2 --- /dev/null +++ b/src/brittle_star_project/evaluation/video.py @@ -0,0 +1,335 @@ +from __future__ import annotations + +import datetime +from pathlib import Path + +import numpy as np +import yaml + +from brittle_star_project import BrittleStarEnv +from brittle_star_project.evaluation.policy import ControlPolicy +from brittle_star_project.evaluation.rollout import ( + EpisodeResult, + _get_observations, + _get_xy_distance_to_target, + _target_reached, + _maybe_clip_action, +) + + +def create_evaluation_dir(model_path: Path) -> Path: + """Create a unique timestamped directory for saving evaluation results.""" + timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + eval_dir = model_path.parent / f"{model_path.stem}_evaluations" / f"eval_{timestamp}" + eval_dir.mkdir(parents=True, exist_ok=True) + return eval_dir + + +def _ensure_offscreen_size(model, width: int, height: int) -> None: + vis_global = getattr(getattr(model, "vis", None), "global_", None) + if vis_global is None: + return + vis_global.offwidth = int(max(width, vis_global.offwidth)) + vis_global.offheight = int(max(height, vis_global.offheight)) + + +def _apply_camera_overrides( + model, + *, + camera_fovy: dict[int, float] | None = None, + camera_xyz: tuple[ + dict[int, float] | None, + dict[int, float] | None, + dict[int, float] | None, + ] = (None, None, None), +) -> None: + if not camera_fovy and not (camera_xyz[0] or camera_xyz[1] or camera_xyz[2]): + return + + ncam = int(getattr(model, "ncam", 0)) + for cam_id, fovy in (camera_fovy or {}).items(): + if cam_id < 0 or cam_id >= ncam: + raise ValueError(f"Camera id {cam_id} is out of range") + model.cam_fovy[cam_id] = float(fovy) + + for cam_id, x in (camera_xyz[0] or {}).items(): + if cam_id < 0 or cam_id >= ncam: + raise ValueError(f"Camera id {cam_id} is out of range") + model.cam_pos[cam_id][0] = float(x) + + for cam_id, y in (camera_xyz[1] or {}).items(): + if cam_id < 0 or cam_id >= ncam: + raise ValueError(f"Camera id {cam_id} is out of range") + model.cam_pos[cam_id][1] = float(y) + + for cam_id, z in (camera_xyz[2] or {}).items(): + if cam_id < 0 or cam_id >= ncam: + raise ValueError(f"Camera id {cam_id} is out of range") + model.cam_pos[cam_id][2] = float(z) + + +def hex_to_rgba(hex_color: str, alpha: float) -> np.ndarray: + color = hex_color.lstrip("#") + if len(color) != 6: + raise ValueError(f"Expected a 6-digit hex color, got {hex_color!r}") + red = int(color[0:2], 16) / 255.0 + green = int(color[2:4], 16) / 255.0 + blue = int(color[4:6], 16) / 255.0 + return np.asarray([red, green, blue, float(alpha)], dtype=np.float32) + + +def save_evaluation_metadata( + eval_dir: Path, + *, + morphology_override_path: str | None, + seed: int, + max_steps: int | None, + result: EpisodeResult, +) -> None: + """Save metadata about the evaluation run.""" + metadata = { + "timestamp": datetime.datetime.now().isoformat(), + "morphology_override": morphology_override_path, + "seed": seed, + "max_steps": max_steps, + "result": { + "return": float(result.return_), + "length": int(result.length), + "reached_target": bool(result.reached_target), + "final_xy_dist": float(result.final_xy_dist) + if result.final_xy_dist is not None + else None, + }, + } + with open(eval_dir / "evaluation_metadata.yaml", "w") as f: + yaml.safe_dump(metadata, f, sort_keys=False) + + +def record_episode( + *, + env: BrittleStarEnv, + policy: ControlPolicy, + seed: int, + max_steps: int, + action_low: np.ndarray | None, + action_high: np.ndarray | None, + action_mask: np.ndarray | None = None, + output_path: Path, + camera_id: int = 1, + fps: int = 60, + width: int = 640, + height: int = 480, + target_xy: tuple[float, float] | None = None, +) -> EpisodeResult: + """Run an episode headlessly and record a video using MuJoCo's Renderer and imageio. + + Args: + env: The environment. + policy: The policy agent. + seed: Random seed. + max_steps: Maximum number of steps. + action_low: Minimum action values. + action_high: Maximum action values. + action_mask: Boolean mask for the actions. + output_path: Where to save the .mp4 file. + camera_id: Camera index to use for rendering (1 is usually close-up). + fps: Frames per second for the video. + width: Video width. + height: Video height. + """ + try: + import imageio + import mujoco + except ImportError as e: + raise ImportError( + "Video recording requires 'imageio' and 'mujoco'. " + "Please install the evaluation dependencies: `uv pip install .[evaluation]`" + ) from e + + state = env.reset(seed=seed, target_position=target_xy) + model = state.mj_model + data = state.mj_data + + _ensure_offscreen_size(model, width, height) + + renderer = mujoco.Renderer(model, width=width, height=height) + ep_return = 0.0 + observations = _get_observations(state) + prev_dist = _get_xy_distance_to_target(observations) if observations else None + initial_dist = prev_dist + reached_target = _target_reached(state=state) + + frames = [] + steps = 0 + + for _ in range(int(max_steps)): + # Capture frame + renderer.update_scene(data, camera=camera_id) + frames.append(renderer.render()) + + # Step environment + obs_dict = observations or {} + action = policy.act(observations=obs_dict) + if action_mask is not None: + action = action[action_mask] + action = _maybe_clip_action(action, action_low, action_high) + + state = env.step(state=state, action=action) + steps += 1 + + observations = _get_observations(state) + cur_dist = _get_xy_distance_to_target(observations) if observations else None + if prev_dist is not None and cur_dist is not None: + ep_return += prev_dist - cur_dist + prev_dist = cur_dist + + reached_target = _target_reached(state=state) + if reached_target: + break + + # Capture final frame + renderer.update_scene(data, camera=camera_id) + frames.append(renderer.render()) + renderer.close() + + # Save video + imageio.mimsave(str(output_path), frames, fps=fps) + + final_dist = _get_xy_distance_to_target(observations) if observations else None + return EpisodeResult( + return_=ep_return, + length=steps, + reached_target=reached_target, + final_xy_dist=final_dist, + initial_target_distance=initial_dist, + ) + + +def record_episode_multi_camera( + *, + env: BrittleStarEnv, + policy: ControlPolicy, + seed: int, + max_steps: int, + action_low: np.ndarray | None, + action_high: np.ndarray | None, + output_paths: dict[int, Path], + action_mask: np.ndarray | None = None, + camera_ids: list[int] | None = None, + camera_fovy: dict[int, float] | None = None, + camera_xyz: tuple[ + dict[int, float] | None, + dict[int, float] | None, + dict[int, float] | None, + ] = (None, None, None), + target_xy: tuple[float, float] | None = None, + robot_color: str | None = None, + fps: int = 60, + width: int = 640, + height: int = 480, +) -> EpisodeResult: + """Run one episode and render multiple camera views to separate files.""" + try: + import imageio + import mujoco + except ImportError as e: + raise ImportError( + "Video recording requires 'imageio' and 'mujoco'. " + "Please install the evaluation dependencies: `uv pip install .[evaluation]`" + ) from e + + if camera_ids is None: + camera_ids = list(output_paths.keys()) + + for cam_id in camera_ids: + if cam_id not in output_paths: + raise ValueError(f"Missing output path for camera {cam_id}") + + output_paths = {cam_id: output_paths[cam_id] for cam_id in camera_ids} + + for path in output_paths.values(): + path.parent.mkdir(parents=True, exist_ok=True) + + state = env.reset(seed=seed, target_position=(target_xy[0], target_xy[1], 0.0)) + model = state.mj_model + data = state.mj_data + + _apply_camera_overrides(model, camera_fovy=camera_fovy, camera_xyz=camera_xyz) + + if robot_color is not None: + robot_body_id = mujoco.mj_name2id( + model, mujoco.mjtObj.mjOBJ_BODY, "BrittleStarMorphology/central_disk" + ) + if robot_body_id < 0: + raise ValueError("Body 'BrittleStarMorphology/central_disk' not found in the model") + + robot_rgba = hex_to_rgba(robot_color, 1.0) + body_parent = model.body_parentid + robot_body_ids = {int(robot_body_id)} + + for body_id in range(1, int(model.nbody)): + current_body_id = int(body_id) + while current_body_id not in (-1, 0, int(robot_body_id)): + current_body_id = int(body_parent[current_body_id]) + if current_body_id == int(robot_body_id): + robot_body_ids.add(body_id) + + for geom_id in range(int(model.ngeom)): + if int(model.geom_bodyid[geom_id]) in robot_body_ids: + model.geom_rgba[geom_id][:] = robot_rgba + + _ensure_offscreen_size(model, width, height) + + renderer = mujoco.Renderer(model, width=width, height=height) + writers = { + cam_id: imageio.get_writer(str(path), fps=fps) for cam_id, path in output_paths.items() + } + + ep_return = 0.0 + observations = _get_observations(state) + prev_dist = _get_xy_distance_to_target(observations) if observations else None + initial_dist = prev_dist + reached_target = _target_reached(state=state) + + steps = 0 + try: + for _ in range(int(max_steps)): + for cam_id in camera_ids: + renderer.update_scene(data, camera=cam_id) + writers[cam_id].append_data(renderer.render()) + + obs_dict = observations or {} + action = policy.act(observations=obs_dict) + if action_mask is not None: + action = action[action_mask] + action = _maybe_clip_action(action, action_low, action_high) + + state = env.step(state=state, action=action) + steps += 1 + + observations = _get_observations(state) + cur_dist = _get_xy_distance_to_target(observations) if observations else None + if prev_dist is not None and cur_dist is not None: + ep_return += prev_dist - cur_dist + prev_dist = cur_dist + + reached_target = _target_reached(state=state) + if reached_target: + break + + for cam_id in camera_ids: + renderer.update_scene(data, camera=cam_id) + writers[cam_id].append_data(renderer.render()) + finally: + renderer.close() + for writer in writers.values(): + writer.close() + + final_dist = _get_xy_distance_to_target(observations) if observations else None + return EpisodeResult( + return_=ep_return, + length=steps, + reached_target=reached_target, + final_xy_dist=final_dist, + initial_target_distance=initial_dist, + ) diff --git a/src/brittle_star_project/ppo.py b/src/brittle_star_project/ppo.py new file mode 100644 index 0000000..fb27ccd --- /dev/null +++ b/src/brittle_star_project/ppo.py @@ -0,0 +1,201 @@ +from functools import partial + +import jax +import jax.numpy as jnp +from jax import debug +from flax.core import FrozenDict +from experiment_logger import get_logger +from brittle_star_project.utils import logged_jit + +logger = get_logger() + + +# Chose to use a class as it seemed the easiest way to integrate the CleanRL code style +# with our need to seperate concerns +class PPO: + def __init__( + self, + args, + sensor_apply, + actor_apply, + critic_apply, + feature_extractor_apply, + message_passer=None, + ): + self.args = args + + if not message_passer: + message_passer = identity + + self.ppo_loss_grad_fn = jax.value_and_grad( + partial( + ppo_loss, + args=args, + sensor_apply=sensor_apply, + actor_apply=actor_apply, + critic_apply=critic_apply, + feature_extractor_apply=feature_extractor_apply, + message_passer=message_passer, + ), + has_aux=True, + ) + + # This PPO class should be initialized only once, + # or this function will need to recompile + @partial(logged_jit, static_argnums=0) + def update_ppo(self, agent_state, storage, key): + debug.callback(logger.debug, f"[PPO] storage.obs shape: {storage.obs.shape}") + debug.callback(logger.debug, f"[PPO] storage.actions shape: {storage.actions.shape}") + debug.callback(logger.debug, f"[PPO] storage.logprobs shape: {storage.logprobs.shape}") + debug.callback(logger.debug, f"[PPO] storage.advantages shape: {storage.advantages.shape}") + debug.callback(logger.debug, f"[PPO] storage.returns shape: {storage.returns.shape}") + + args = self.args + ppo_loss_grad_fn = self.ppo_loss_grad_fn + + def update_epoch(carry, _): + agent_state, key = carry + key, subkey = jax.random.split(key) + + def flatten(x): + return x.reshape((-1,) + x.shape[2:]) + + def convert_data(x): + x = jax.random.permutation(subkey, x) + return jnp.reshape(x, (args.num_minibatches, -1) + x.shape[1:]) + + flatten_storage = jax.tree.map(flatten, storage) + shuffled_storage = jax.tree.map(convert_data, flatten_storage) + + def update_minibatch(agent_state, minibatch): + debug.callback(logger.debug, f"[PPO] minibatch.obs: {minibatch.obs.shape}") + debug.callback(logger.debug, f"[PPO] minibatch.actions: {minibatch.actions.shape}") + debug.callback( + logger.debug, f"[PPO] minibatch.logprobs: {minibatch.logprobs.shape}" + ) + debug.callback( + logger.debug, f"[PPO] minibatch.advantages: {minibatch.advantages.shape}" + ) + debug.callback(logger.debug, f"[PPO] minibatch.returns: {minibatch.returns.shape}") + + (loss, (pg_loss, v_loss, entropy_loss, approx_kl)), grads = ppo_loss_grad_fn( + agent_state.params, + minibatch.obs, + minibatch.actions, + minibatch.logprobs, + minibatch.advantages, + minibatch.returns, + ) + agent_state = agent_state.apply_gradients(grads=grads) + return agent_state, (loss, pg_loss, v_loss, entropy_loss, approx_kl) + + agent_state, metrics = jax.lax.scan(update_minibatch, agent_state, shuffled_storage) + return (agent_state, key), metrics + + (agent_state, key), (loss, pg_loss, v_loss, entropy_loss, approx_kl) = jax.lax.scan( + update_epoch, (agent_state, key), (), length=args.update_epochs + ) + return agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, key + + +""" +Should be ok to use partial here, since the references to network, +actor and critic should not change at runtime +The cost of seperating concerns is to somehow pass these values +that are now not in the same scope +""" + + +@partial(logged_jit, static_argnums=(0, 1, 2, 3, 4)) +def get_action_and_value( + sensor_apply, + actor_apply, + message_passer, + critic_apply, + feature_extractor_apply, + params: FrozenDict, + x: jnp.ndarray, + action: jnp.ndarray, +): + hidden_sensor = sensor_apply(params["sensor_params"], x) + hidden_critic = feature_extractor_apply(params["feature_extractor_params"], x) + + # only apply message passing in decentralized context + if message_passer is not None: + hidden_sensor = message_passer(params["message_passer_params"], hidden_sensor) + + debug.callback(logger.debug, f"[SHAPE] hidden_sensor: {hidden_sensor.shape}") + debug.callback(logger.debug, f"[SHAPE] hidden_critic: {hidden_critic.shape}") + + mean, log_std = actor_apply(params["actor_params"], hidden_sensor) + + debug.callback(logger.debug, f"[SHAPE] mean: {mean.shape}") + debug.callback(logger.debug, f"[SHAPE] log_std: {log_std.shape}") + debug.callback(logger.debug, f"[SHAPE] action: {action.shape}") + + log_std = jnp.clip(log_std, -5, 2) + std = jnp.exp(log_std) + + logprob = -0.5 * (((action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)) + debug.callback(logger.debug, f"[SHAPE] logprob pre-sum: {logprob.shape}") + + logprob = logprob.sum(axis=(-2, -1)) + debug.callback(logger.debug, f"[SHAPE] logprob final: {logprob.shape}") + + entropy = (0.5 + 0.5 * jnp.log(2 * jnp.pi) + log_std).sum(axis=(-2, -1)) + value = critic_apply(params["critic_params"], hidden_critic).squeeze(-1) + debug.callback(logger.debug, f"[SHAPE] value: {value.shape}") + + return logprob, entropy, value + + +def ppo_loss( + params, + x, + a, + logp, + mb_advantages, + mb_returns, + args, + sensor_apply, + actor_apply, + message_passer, + critic_apply, + feature_extractor_apply, +): + newlogprob, entropy, newvalue = get_action_and_value( + sensor_apply, + actor_apply, + message_passer, + critic_apply, + feature_extractor_apply, + params, + x, + a, + ) + logratio = newlogprob - logp + ratio = jnp.exp(logratio) + approx_kl = ((ratio - 1) - logratio).mean() + + if args.norm_adv: + mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8) + + pg_loss1 = -mb_advantages * ratio + pg_loss2 = -mb_advantages * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef) + pg_loss = jnp.maximum(pg_loss1, pg_loss2).mean() + + v_loss = 0.5 * ((newvalue - mb_returns) ** 2).mean() + entropy_loss = entropy.mean() + loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef + return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl)) + + +def identity(_, hidden): + """ + Used for seamless jax integration, + avoids having branching inside jitted function, + used as message_passer in case it is not given, + (in case of centralized lvl) + """ + + return hidden diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py new file mode 100644 index 0000000..98d4045 --- /dev/null +++ b/src/brittle_star_project/trainers/PPOTrainer.py @@ -0,0 +1,988 @@ +import datetime +import random +import time +from dataclasses import asdict, dataclass +from functools import partial +from typing import Any, Optional + +import jax +import jax.numpy as jnp +import numpy as np +import optax +import flax.linen as nn +from flax.training.train_state import TrainState + +from experiment_logger import get_logger + +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.dataclasses import EpisodeStatistics +from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper +from brittle_star_project.environment.obs_processing import create_obs_processor +from brittle_star_project.evaluation.evaluate_mjx import ( + append_checkpoint_eval_row, + build_eval_rollout_fn, + evaluate_checkpoint_mjx, +) +from brittle_star_project.MLPs.routing import apply_per_node +from brittle_star_project.MLPs.mlps import ( + Actor, + AgentParams, + GenericDenseLayersWithActivation, + MessagePasser, + OneDenseLayerMLP, + Storage, +) +from brittle_star_project.MLPs.adjancency_builder import build_adjacency +from brittle_star_project.ppo import PPO +from brittle_star_project.environment import MorphMode +from brittle_star_project.utils import logged_jit + +from brittle_star_project.environment.env_types import Backend + +# TODO: clip scaled reward? + + +@logged_jit +def _clip_action(action: jnp.ndarray, low: jnp.ndarray, high: jnp.ndarray) -> jnp.ndarray: + return jnp.clip(action, low, high) + + +def _compute_explained_variance(values: jnp.ndarray, returns: jnp.ndarray) -> float: + var_returns = jnp.var(returns) + explained_var = 1.0 - jnp.var(returns - values) / (var_returns + 1e-8) + return float(explained_var) + + +@logged_jit +def _linear_schedule(count, minibatch_count, update_epochs, num_iterations, learning_rate): + frac = 1.0 - (count // (minibatch_count * update_epochs)) / num_iterations + return learning_rate * frac + + +def _get_action_and_value_noise( + sensor: nn.Module, + feature_extractor: nn.Module, + actor: nn.Module, + critic: nn.Module, + message_passer: Optional[nn.Module], + agent_state: TrainState, + next_obs: jnp.ndarray, + key, + action_low, + action_high, +): + # (B, n_nodes, feat) + hidden = apply_per_node(sensor.apply, agent_state.params["sensor_params"], next_obs) + + if message_passer is not None: + params = agent_state.params["message_passer_params"] + # (n_nodes, feat) --> let each node talk with its neighbours ==> vmap over B dimension + hidden = jax.vmap(lambda x: message_passer.apply(params, x))(hidden) + + hidden_critic = apply_shared( + feature_extractor, agent_state.params["feature_extractor_params"], next_obs + ) + + mean, log_std = apply_per_node(actor.apply, agent_state.params["actor_params"], hidden) + log_std = jnp.clip(log_std, -5, 2) + key, subkey = jax.random.split(key) + noise = jax.random.normal(subkey, shape=mean.shape) + std = jnp.exp(log_std) + + raw_action = mean + noise * std + flat_action = raw_action.reshape( + raw_action.shape[0], -1 + ) # concat the per agent, keep the envs dim (batch, agent * action) + flat_clipped_action = _clip_action(flat_action, action_low, action_high) + + logprob = -0.5 * (((raw_action - mean) / std) ** 2 + 2 * log_std + jnp.log(2 * jnp.pi)).sum( + axis=(-2, -1) + ) + value = apply_shared(critic, agent_state.params["critic_params"], hidden_critic) + + return flat_clipped_action, raw_action, logprob, value.squeeze(-1), mean, std, key + + +def _step_once( + carry, + _, + env_step_fn, + num_envs: int, + sensor: nn.Module, + feature_extractor: nn.Module, + actor: nn.Module, + critic: nn.Module, + message_passer: Optional[nn.Module], + action_low, + action_high, +): + agent_state, episode_stats, obs, done, key, env_state, terminated_any, truncated_any = carry + flat_clipped_action, raw_action, logprob, value, mean, std, key = _get_action_and_value_noise( + sensor, + feature_extractor, + actor, + critic, + message_passer, + agent_state, + obs, + key, + action_low, + action_high, + ) + logger = get_logger() + + logger.debug(f"[_step_once] raw_action: {raw_action.shape}") + logger.debug(f"[_step_once] clipped_action: {flat_clipped_action.shape}") + + # Supporting signals (often where mismatch originates) + logger.debug(f"[_step_once] logprob: {logprob.shape}") + logger.debug(f"[_step_once] value: {value.shape}") + logger.debug(f"[_step_once] mean: {mean.shape}") + logger.debug(f"[_step_once] std: {std.shape}") + + key, reset_key = jax.random.split(key) + reset_rngs = jax.random.split(reset_key, num_envs) + + # ---- ENV STEP ---- + key, reset_key = jax.random.split(key) + reset_rngs = jax.random.split(reset_key, num_envs) + + episode_stats, env_state, (next_obs, reward, next_done, terminated, truncated) = env_step_fn( + episode_stats, + env_state, + flat_clipped_action, + reset_rngs, + ) + + terminated_any = terminated_any | terminated + truncated_any = truncated_any | truncated + + logger.debug(f"[_step_once] next_obs: {next_obs.shape}") + logger.debug(f"[_step_once] reward: {reward.shape}") + logger.debug(f"[_step_once] next_done: {next_done.shape}") + + storage = Storage( + obs=obs, + actions=raw_action, + raw_actions=raw_action, + logprobs=logprob, + dones=done, + values=value, + rewards=reward, + means=mean, + stds=std, + returns=jnp.zeros_like(reward), + advantages=jnp.zeros_like(reward), + ) + return ( + agent_state, + episode_stats, + next_obs, + next_done, + key, + env_state, + terminated_any, + truncated_any, + ), storage + + +def reward_fn(env_state, next_env_state): + """Shaped reward used during training and checkpoint evaluation. + + Public so that ``evaluation.evaluate_mjx`` can import it and produce + metrics that are directly comparable to training-time returns. + """ + # Positive delta_distance means the brittle star is moving *away* from target. + delta_distance = ( + next_env_state.observations["xy_distance_to_target"] + - env_state.observations["xy_distance_to_target"] + ).squeeze(-1) + + env_reward = next_env_state.reward + clipped_env_reward = jnp.clip(100 * env_reward, -10, 10) + + time_penalty = 0.1 + distance_penalty = jnp.clip(0.5 * delta_distance, -0.5, 0.5) + penalty = time_penalty + distance_penalty + + return jnp.where(next_env_state.terminated, 50.0, clipped_env_reward - penalty) + + +def _step_env_wrapped( + episode_stats, + env_state, + action, + reset_rngs, + env_step_fn, + reset_single_fn, + obs_processor, +): + next_env_state_pre_reset = env_step_fn(env_state, action) + + reward = reward_fn(env_state, next_env_state_pre_reset) + terminated = next_env_state_pre_reset.terminated + truncated = next_env_state_pre_reset.truncated + done = terminated | truncated + + new_episode_return = episode_stats.episode_returns + reward + new_episode_length = episode_stats.episode_lengths + 1 + + episode_stats = episode_stats.replace( + episode_returns=new_episode_return * (1 - done), + episode_lengths=new_episode_length * (1 - done), + returned_episode_returns=jnp.where( + done, new_episode_return, episode_stats.returned_episode_returns + ), + returned_episode_lengths=jnp.where( + done, new_episode_length, episode_stats.returned_episode_lengths + ), + ) + + def _maybe_reset(state_i, rng_i, do_reset_i): + def _do(_): + reset_state = reset_single_fn(rng=rng_i) + + def _cast_leaf(new_leaf, like_leaf): + if like_leaf is None or new_leaf is None: + return new_leaf + + # Use jnp.asarray(...) to robustly get dtype for both JAX arrays and Python scalars. + like_dtype = jnp.asarray(like_leaf).dtype + + # Avoid unnecessary work when already matching. + if hasattr(new_leaf, "dtype") and new_leaf.dtype == like_dtype: + return new_leaf + + return jnp.asarray(new_leaf, dtype=like_dtype) + + # `lax.cond` requires both branches to return identical PyTree types/dtypes. + return jax.tree_util.tree_map(_cast_leaf, reset_state, state_i) + + def _dont(_): + return state_i + + return jax.lax.cond(do_reset_i, _do, _dont, operand=None) + + # Auto-reset done envs so rollouts continue with fresh episode initial states. + next_env_state = jax.vmap(_maybe_reset)(next_env_state_pre_reset, reset_rngs, done) + + return ( + episode_stats, + next_env_state, + (obs_processor(next_env_state.observations), reward, done, terminated, truncated), + ) + + +def apply_shared(net, params, x): + # x: (batch, nodes, feat) + # If the critic expects a single vector per environment: + batch_size = x.shape[0] + x_flattened = x.reshape(batch_size, -1) + return jax.vmap(lambda xi: net.apply(params, xi))(x_flattened) + + +def _rollout_jit( + agent_state, + episode_stats, + env_state, + next_obs, + next_done, + key, + max_steps, + step_env_fn, + num_envs: int, + sensor: nn.Module, + feature_extractor: nn.Module, + actor: nn.Module, + critic: nn.Module, + message_passer: Optional[nn.Module], + action_low, + action_high, +): + terminated_any0 = jnp.zeros((num_envs,), dtype=jnp.bool_) + truncated_any0 = jnp.zeros((num_envs,), dtype=jnp.bool_) + + ( + ( + agent_state, + episode_stats, + next_obs, + next_done, + key, + env_state, + terminated_any, + truncated_any, + ), + storage, + ) = jax.lax.scan( + partial( + _step_once, + sensor=sensor, + feature_extractor=feature_extractor, + actor=actor, + critic=critic, + message_passer=message_passer, + env_step_fn=step_env_fn, + num_envs=num_envs, + action_low=action_low, + action_high=action_high, + ), + ( + agent_state, + episode_stats, + next_obs, + next_done, + key, + env_state, + terminated_any0, + truncated_any0, + ), + (), + max_steps, + ) + return ( + agent_state, + episode_stats, + next_obs, + next_done, + storage, + key, + env_state, + terminated_any, + truncated_any, + ) + + +def _compute_gae_once(carry, inp, gamma, gae_lambda): + advantages = carry + nextdone, nextvalues, curvalues, reward = inp + nextnonterminal = 1.0 - nextdone + delta = reward + gamma * nextvalues * nextnonterminal - curvalues + advantages = delta + gamma * gae_lambda * nextnonterminal * advantages + return advantages, advantages + + +def _compute_gae_jit( + agent_state, + storage, + next_obs, + next_done, + gamma, + gae_lambda, + num_envs, + feature_extractor, + critic, +): + next_value = apply_shared( + critic, + agent_state.params["critic_params"], + apply_shared(feature_extractor, agent_state.params["feature_extractor_params"], next_obs), + ).squeeze(-1) + + advantages = jnp.zeros((num_envs,)) + dones = jnp.concatenate([storage.dones, next_done[None, :]], axis=0) + values = jnp.concatenate([storage.values, next_value[None, :]], axis=0) + _, advantages = jax.lax.scan( + partial(_compute_gae_once, gamma=gamma, gae_lambda=gae_lambda), + advantages, + (dones[1:], values[1:], values[:-1], storage.rewards), + reverse=True, + ) + returns = advantages + storage.values + advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8) + return storage.replace(advantages=advantages, returns=returns) + + +@dataclass +class TrainingMeasurements: + loss: jnp.ndarray + pg_loss: jnp.ndarray + v_loss: jnp.ndarray + entropy_loss: jnp.ndarray + approx_kl: jnp.ndarray + avg_episodic_return: float + explained_variance: float + num_terminated: int + num_truncated: int + avg_terminated_length: Any + avg_truncated_length: Any + + +class PPOTrainer: + def __init__( + self, + cfg: BrittleStarConfig, + env: BrittleStarJaxEnvWrapper, + run_dir: str, + run_name: str, + ): + self.cfg = cfg + self.ppo = cfg.ppo + self.experiment = cfg.experiment + self.logging_cfg = cfg.logging + self.evaluation_cfg = cfg.evaluation + self.env = env + self.run_dir = run_dir + self.run_name = run_name + self.logger = get_logger() + + # Derived runtime fields + self.batch_size = self.ppo.num_envs * self.ppo.num_steps + self.num_iterations = self.ppo.total_timesteps // self.batch_size + + self.key = jax.random.PRNGKey(self.experiment.seed) + + self.morph_mode = self.cfg.morphology.morph_mode + + self.segments_per_arm = jnp.asarray(self.cfg.morphology.segments_per_arm, dtype=jnp.int32) + self.num_segments = self.segments_per_arm.sum().item() + self.num_arms = jnp.where(self.segments_per_arm > 0, 1, 0).sum().item() + + self.logger.info(f"[INIT]: Used morphology mode {self.morph_mode}") + self.adj = build_adjacency(cfg.morphology.segments_per_arm, self.morph_mode) + + ( + self.sensor, + self.message_passer, + self.actor, + self.feature_extractor, + self.critic, + self.needed_copies, + self.agent_indices, + ) = self._init_agent() + + self.sensor.apply = logged_jit(self.sensor.apply) + self.feature_extractor.apply = logged_jit(self.feature_extractor.apply) + self.actor.apply = logged_jit(self.actor.apply) + self.critic.apply = logged_jit(self.critic.apply) + + # Build the centralized observation processor: derive -> normalize -> pad -> flatten. + self.obs_processor = create_obs_processor( + bounds_dict=self.cfg.obs_bounds.to_bounds_dict(), + needed_copies=self.needed_copies, + num_arms=self.num_arms, + morph_mode=self.morph_mode, + padding_masks=self.env.padding_masks, + segments_per_arm=self.segments_per_arm, + agent_indices=self.agent_indices, + ) + + self.logger.debug(f"needed copies = {self.needed_copies}") + + action_low = jnp.asarray(self.env.single_action_space.low, dtype=jnp.float32) + action_high = jnp.asarray(self.env.single_action_space.high, dtype=jnp.float32) + self._action_low = action_low + self._action_high = action_high + + self._rollout_jit = logged_jit( + partial( + _rollout_jit, + max_steps=self.ppo.num_steps, + step_env_fn=partial( + _step_env_wrapped, + env_step_fn=self.env.step, + reset_single_fn=self.env.raw.reset, + obs_processor=self.obs_processor, + ), + num_envs=self.ppo.num_envs, + sensor=self.sensor, + feature_extractor=self.feature_extractor, + actor=self.actor, + critic=self.critic, + message_passer=self.message_passer, + action_low=action_low, + action_high=action_high, + ) + ) + self._compute_gae_jit = logged_jit( + partial( + _compute_gae_jit, + num_envs=self.ppo.num_envs, + gamma=self.ppo.gamma, + gae_lambda=self.ppo.gae_lambda, + feature_extractor=self.feature_extractor, + critic=self.critic, + ) + ) + + def apply_sensor(p, x): + return apply_per_node(self.sensor.apply, p, x) + + def apply_actor(p, x): + return apply_per_node(self.actor.apply, p, x) + + def apply_critic(p, x): + return apply_shared(self.critic, p, x) + + def apply_feature(p, x): + return apply_shared(self.feature_extractor, p, x) + + def apply_message_passer(p, x): + assert self.message_passer is not None + return jax.vmap(lambda x_in: self.message_passer.apply(p, x_in))(x) + + self._ppo = PPO( + self.ppo, + apply_sensor, + apply_actor, + apply_critic, + apply_feature, + apply_message_passer if self.message_passer is not None else None, + ) + + self.agent_state = self._init_agent_state() + + self.episode_stats = self._init_episode_stats() + + self._init_random() + # Lazily-built JIT-compiled MJX eval rollout, created on first evaluation. + self._eval_fn = None + + def _init_random(self): + self.logger.info(f"[RANDOM]: Setting random seed to {self.experiment.seed}") + + random.seed(self.experiment.seed) + np.random.seed(self.experiment.seed) + + def _init_agent(self): + self.logger.info("[AGENT]: Initializing agent...") + agent_indices = [0, 1, 2, 3, 4] + match self.morph_mode: + case MorphMode.CENTRALIZED: + needed_copies = 1 + case MorphMode.FULLY_CONNECTED | MorphMode.RING: + agent_mask = self.segments_per_arm > 0 + agent_indices = jnp.where(agent_mask)[0] + needed_copies = jnp.where(self.segments_per_arm > 0, 1, 0).sum().item() + case MorphMode.SEGMENT: + agent_mask = self.segments_per_arm > 0 + agent_indices = jnp.where(agent_mask)[0] + needed_copies = ( + self.segments_per_arm.sum() + jnp.where(self.segments_per_arm > 0, 1, 0).sum() + ).item() + + # scale actor output with size of model --> more models ==> less actions needed per model + actor = Actor(action_dim=self.env.single_action_space.shape[0] // needed_copies) + sensor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300]) + message_passer: Optional[nn.Module] = ( + MessagePasser( + hidden_dim=300, + num_propagation_steps=self.cfg.architecture.message_passing_steps or 4, + adj_matrix=self.adj, + ) + if self.morph_mode != MorphMode.CENTRALIZED + else None + ) + + feature_extractor = GenericDenseLayersWithActivation(layer_sizes=[300, 300, 300]) + critic = OneDenseLayerMLP() + return ( + sensor, + message_passer, + actor, + feature_extractor, + critic, + needed_copies, + agent_indices, + ) + + def _init_agent_state(self) -> TrainState: + self.logger.info("[AGENT STATE]: Initializing agent state...") + + self.key, sensor_key, actor_key, critic_key, feature_extractor_key, message_passer_key = ( + jax.random.split(self.key, 6) + ) + + dummy_reset = self.env.reset(seed=0) + + for k, v in dummy_reset.observations.items(): + self.logger.debug(k, v.shape) + + sample_obs = self.obs_processor(dummy_reset.observations)[0] # take first env + + self.logger.debug(f"[_init_agent_state] sample_obs: {sample_obs.shape}") + self.obs_mean = jnp.zeros((sample_obs.shape[-1],)) + self.obs_var = jnp.ones((sample_obs.shape[-1],)) + self.obs_count = 1e-4 + self.logger.debug(f"[_init_agent_state] obs_mean: {self.obs_mean.shape}") + self.logger.debug(f"[_init_agent_state] obs_var: {self.obs_var.shape}") + + self.logger.debug(f"[_init_agent_state]: Needed copies: {self.needed_copies}") + sensor_keys = jax.random.split(sensor_key, self.needed_copies) + actor_keys = jax.random.split(actor_key, self.needed_copies) + + # (needed_copies, X) + sensor_params = jax.vmap(lambda k: self.sensor.init(k, sample_obs))(sensor_keys) + self.logger.debug( + f"[_init_agent_state] sensor_params: {jax.tree.map(lambda x: x.shape, sensor_params)}" + ) + + single_sensor_param = jax.tree.map(lambda x: x[0], sensor_params) + self.logger.debug( + f"[_init_agent_state] single_sensor_param: { + jax.tree.map(lambda x: x.shape, single_sensor_param) + }" + ) + + sensor_params_sample = self.sensor.apply(single_sensor_param, sample_obs) + self.logger.debug( + f"[_init_agent_state] sensor_params_sample shape: {sensor_params_sample.shape}" + ) + + actor_params = jax.vmap(lambda k: self.actor.init(k, sensor_params_sample))(actor_keys) + self.logger.debug( + f"[_init_agent_state] actor_params: {jax.tree.map(lambda x: x.shape, actor_params)}" + ) + + message_passer_params = {} + if self.morph_mode != MorphMode.CENTRALIZED: + assert self.message_passer is not None, "decentralized modes require a message passer" + + message_passer_params = self.message_passer.init( + message_passer_key, + self.sensor.apply(single_sensor_param, sample_obs), + ) + self.logger.debug( + f"[_init_agent_state] message_passer_params: { + jax.tree.map(lambda x: x.shape, message_passer_params) + }" + ) + + flat_obs = sample_obs.reshape(-1) # BECAUSE 1 centralized critic + self.logger.debug(f"[_init_agent_state] flat_obs: {flat_obs.shape}") + + feature_extractor_params = self.feature_extractor.init(feature_extractor_key, flat_obs) + self.logger.debug( + f"[_init_agent_state] feature_extractor_params: { + jax.tree.map(lambda x: x.shape, feature_extractor_params) + }" + ) + + critic_input = self.feature_extractor.apply(feature_extractor_params, flat_obs) + self.logger.debug(f"[_init_agent_state] critic_input: {critic_input.shape}") + + critic_params = self.critic.init(critic_key, critic_input) + self.logger.debug( + f"[_init_agent_state] critic_params: {jax.tree.map(lambda x: x.shape, critic_params)}" + ) + + return TrainState.create( + apply_fn=None, + params=asdict( + AgentParams( + sensor_params, + actor_params, + critic_params, + feature_extractor_params, + message_passer_params, + ) + ), + tx=optax.chain( + optax.clip_by_global_norm(self.ppo.max_grad_norm), + optax.inject_hyperparams(optax.adam)( + learning_rate=partial( + _linear_schedule, + minibatch_count=self.ppo.num_minibatches, + update_epochs=self.ppo.update_epochs, + num_iterations=self.num_iterations, + learning_rate=self.ppo.learning_rate, + ) + if self.ppo.anneal_lr + else self.ppo.learning_rate, + eps=1e-5, + ), + ), + ) + + def _init_episode_stats(self) -> EpisodeStatistics: + self.logger.info("[EPISODE STATS]: Initializing episode stats...") + + return EpisodeStatistics( + episode_returns=jnp.zeros(self.ppo.num_envs, dtype=jnp.float32), + episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32), + returned_episode_returns=jnp.zeros(self.ppo.num_envs, jnp.float32), + returned_episode_lengths=jnp.zeros(self.ppo.num_envs, dtype=jnp.int32), + ) + + def _rollout(self, env_state, next_obs, next_done) -> tuple[Any, ...]: + return self._rollout_jit( + self.agent_state, + self.episode_stats, + env_state, + next_obs, + next_done, + self.key, + ) + + def _compute_gae(self, storage, next_obs, next_done) -> Storage: + return self._compute_gae_jit( + self.agent_state, + storage, + next_obs, + next_done, + ) + + def _log( + self, + global_step, + episode_stats, + start_time, + iteration_time_start, + training_measurements, + storage, + ): + data = jax.device_get( + { + "rewards": storage.rewards, + "values": storage.values, + "returns": storage.returns, + "advantages": storage.advantages, + } + ) + + rollout_metrics = { + "rollout/reward_mean": float(np.mean(data["rewards"])), + "rollout/return_mean": float(np.mean(data["returns"])), + "rollout/value_mean": float(np.mean(data["values"])), + "rollout/advantage_mean": float(np.mean(data["advantages"])), + "rollout/advantage_std": float(np.std(data["advantages"])), + "rollout/value_vs_return_mse": float(np.mean((data["values"] - data["returns"]) ** 2)), + } + + metrics = { + "charts/episodic_return": training_measurements.avg_episodic_return, + "charts/episodic_length": float( + np.mean(jax.device_get(episode_stats.returned_episode_lengths)) + ), + "charts/explained_variance": training_measurements.explained_variance, + "losses/value_loss": training_measurements.v_loss[-1, -1].item(), + "losses/policy_loss": training_measurements.pg_loss[-1, -1].item(), + "losses/entropy": training_measurements.entropy_loss[-1, -1].item(), + "losses/approx_kl": training_measurements.approx_kl[-1, -1].item(), + "charts/learning_rate": self.agent_state.opt_state[1] + .hyperparams["learning_rate"] + .item(), + "charts/SPS": int(global_step / (time.time() - start_time)), + "charts/SPS_update": int( + self.ppo.num_envs * self.ppo.num_steps / (time.time() - iteration_time_start) + ), + "termi_trunci/num_terminated": training_measurements.num_terminated, + "termi_trunci/num_truncated": training_measurements.num_truncated, + "termi_trunci/avg_terminated_ep_length": training_measurements.avg_terminated_length, + "termi_trunci/avg_truncated_ep_length": training_measurements.avg_truncated_length, + **rollout_metrics, + } + + self.logger.log(metrics, step=global_step) + + def _step(self, env_state, next_obs, next_done, iteration: int) -> tuple: + if iteration == 1: + self.logger.log_non_interactive(f"Starting first rollout (JIT): {time.ctime()}") + self.logger.debug(f"[_step] next_obs (in): {next_obs.shape}") + ( + self.agent_state, + self.episode_stats, + next_obs, + next_done, + storage, + self.key, + next_env_state, + terminated_any, + truncated_any, + ) = self._rollout(env_state, next_obs, next_done) + self.logger.debug(f"[_step] next_obs (post-rollout): {next_obs.shape}") + if iteration == 1: + self.logger.log_non_interactive(f"First rollout completed: {time.ctime()}") + + storage = self._compute_gae(storage, next_obs, next_done) + self.logger.debug(f"[_step] storage.obs (post-gae): {storage.obs.shape}") + if iteration == 1: + self.logger.log_non_interactive(f"Starting first PPO update (JIT): {time.ctime()}") + + self.agent_state, loss, pg_loss, v_loss, entropy_loss, approx_kl, self.key = ( + self._ppo.update_ppo(self.agent_state, storage, self.key) + ) + + if iteration == 1: + self.logger.log_non_interactive(f"First PPO update completed: {time.ctime()}") + + avg_episodic_return = float( + jnp.mean(jax.device_get(self.episode_stats.returned_episode_returns)).item() + ) + + explained_var = _compute_explained_variance(storage.values, storage.returns) + + terminated = terminated_any + truncated = truncated_any + episode_lengths = self.episode_stats.returned_episode_lengths + + num_terminated = int(jnp.sum(terminated).item()) + num_truncated = int(jnp.sum(truncated).item()) + + avg_terminated_length = jnp.sum(episode_lengths * terminated) / jnp.maximum( + jnp.sum(terminated), 1 + ) + + avg_truncated_length = jnp.sum(episode_lengths * truncated) / jnp.maximum( + jnp.sum(truncated), 1 + ) + + return ( + next_env_state, + next_obs, + next_done, + TrainingMeasurements( + loss=loss, + pg_loss=pg_loss, + v_loss=v_loss, + entropy_loss=entropy_loss, + approx_kl=approx_kl, + avg_episodic_return=avg_episodic_return, + explained_variance=explained_var, + num_terminated=num_terminated, + num_truncated=num_truncated, + avg_terminated_length=avg_terminated_length, + avg_truncated_length=avg_truncated_length, + ), + storage, + ) + + def _close(self): + self.env.close() + + def _save_model(self, model_path: str): + self.logger.info("[SAVE]: Saving the final model...") + self.logger.save_final_model(params=self.agent_state.params, metadata=asdict(self.cfg)) + + def _save_checkpoint(self, iteration: int): + self.logger.info(f"[SAVE]: Saving checkpoint at iteration {iteration}...") + self.logger.save_checkpoint( + params=self.agent_state.params, step=iteration, metadata=asdict(self.cfg) + ) + + def _evaluate_checkpoint(self, iteration: int, *, trained_timesteps: int) -> None: + """Evaluate the current checkpoint and persist metrics to CSV. + + Delegates all evaluation logic to `evaluation.evaluate_mjx`. + Best-effort: a failure here must never abort training. + """ + if not self.evaluation_cfg.evaluate_checkpoints: + return + + max_steps = int(self.evaluation_cfg.eval_max_steps) + seed = int(self.evaluation_cfg.eval_seed) + + if max_steps <= 0: + self.logger.warning("[EVAL]: eval_max_steps must be > 0; skipping evaluation") + return + + if not self.logging_cfg.save_checkpoints or self.logging_cfg.checkpoint_frequency <= 0: + self.logger.warning( + "[EVAL]: evaluate_checkpoints is enabled but checkpoint saving is disabled; " + "skipping evaluation" + ) + return + + try: + if self._eval_fn is None: + if getattr(self.env, "backend", None) != Backend.MJX: + self.logger.warning( + f"[EVAL]: Training env backend is {self.env.backend}; " + "MJX evaluation may be unavailable/slow." + ) + self._eval_fn = build_eval_rollout_fn( + env=self.env, + obs_processor=self.obs_processor, + sensor_apply=lambda p, x: apply_per_node(self.sensor.apply, p, x), + actor_apply=lambda p, x: apply_per_node(self.actor.apply, p, x), + message_passer_apply=( + None if self.message_passer is None else self.message_passer.apply + ), + action_low=self._action_low, + action_high=self._action_high, + reward_fn=reward_fn, + ) + + result = evaluate_checkpoint_mjx( + self._eval_fn, + self.agent_state.params, + seed=seed, + max_steps=max_steps, + ) + csv_path = append_checkpoint_eval_row( + self.run_dir, + iteration=iteration, + trained_timesteps=int(trained_timesteps), + result=result, + ) + self.logger.sync_file(csv_path) + except Exception as e: + self.logger.warning(f"[EVAL]: Checkpoint evaluation failed: {e}") + + def train(self): + """ + Train the PPO agent for a specified number of iterations. + Closes the environment at the end of training. + """ + self.logger.info(f"running name: {self.run_name}") + + self.logger.info("[TRAIN]: Resetting environment...") + self.logger.log_non_interactive(f"Initial reset started: {time.ctime()}") + + env_state = self.env.reset(seed=self.experiment.seed) + + next_obs = self.obs_processor(env_state.observations) + self.logger.debug(f"[train] next_obs: {next_obs.shape}") + + next_done = jnp.zeros(self.ppo.num_envs, dtype=jnp.bool_) + + self.logger.log_non_interactive(f"Initial reset completed: {time.ctime()}") + + global_step = 0 + start_time = time.time() + + iter_bar = self.logger.progress_bar(range(1, self.num_iterations + 1)) + for iteration in iter_bar: + iteration_time_start = time.time() + + env_state, next_obs, next_done, training_measurements, storage = self._step( + env_state, next_obs, next_done, iteration=iteration + ) + + global_step += self.ppo.num_steps * self.ppo.num_envs + self._log( + global_step, + self.episode_stats, + start_time, + iteration_time_start, + training_measurements, + storage, + ) + + sps = int(global_step / (time.time() - start_time)) + remaining_steps = self.ppo.total_timesteps - global_step + eta_seconds = int(remaining_steps / sps) if sps > 0 else 0 + eta_str = str(datetime.timedelta(seconds=eta_seconds)) + + self.logger.log_non_interactive( + f"Iteration {iteration}/{self.num_iterations} | " + f"Step {global_step}/{self.ppo.total_timesteps} | " + f"SPS {sps} | " + f"Return {training_measurements.avg_episodic_return:.4f} | " + f"ETA {eta_str}" + ) + + if self.logging_cfg.save_checkpoints and self.logging_cfg.checkpoint_frequency > 0: + if iteration % self.logging_cfg.checkpoint_frequency == 0: + self._save_checkpoint(iteration) + self._evaluate_checkpoint(iteration, trained_timesteps=global_step) + + if getattr(self.cfg.experiment, "debug_sanity", False): + self.logger.info("\n[SANITY CHECK] Successfully completed 1 epoch") + break + + if self.logging_cfg.save_model: + model_path = f"{self.run_dir}/{self.experiment.exp_name}.cleanrl_model" + self._save_model(model_path=model_path) + + self._close() diff --git a/configs/.gitkeep b/src/brittle_star_project/trainers/__init__.py similarity index 100% rename from configs/.gitkeep rename to src/brittle_star_project/trainers/__init__.py diff --git a/src/brittle_star_project/utils/__init__.py b/src/brittle_star_project/utils/__init__.py new file mode 100644 index 0000000..ed72e5a --- /dev/null +++ b/src/brittle_star_project/utils/__init__.py @@ -0,0 +1,3 @@ +from .logged_jit import logged_jit + +__all__ = ["logged_jit"] diff --git a/src/brittle_star_project/utils/logged_jit.py b/src/brittle_star_project/utils/logged_jit.py new file mode 100644 index 0000000..3d29ba6 --- /dev/null +++ b/src/brittle_star_project/utils/logged_jit.py @@ -0,0 +1,17 @@ +import jax +from experiment_logger import get_logger + + +def logged_jit(fn, **jit_kwargs): + logger = get_logger() + name = getattr(fn, "__name__", getattr(fn, "__qualname__", repr(fn))) + + def decorator(func): + def traced_func(*args, **kwargs): + logger.debug(f"[JIT] Compiling {name}...") + return func(*args, **kwargs) + + jitted = jax.jit(traced_func, **jit_kwargs) + return jitted + + return decorator(fn) diff --git a/src/experiment_logger/README.md b/src/experiment_logger/README.md new file mode 100644 index 0000000..4e13b6b --- /dev/null +++ b/src/experiment_logger/README.md @@ -0,0 +1,71 @@ +# Experiment Logger + +A standardized, unified interface for logging experiments across multiple backends (WandB, TensorBoard, and Local Disk). + +This library is designed to be a standalone package that decouples the logging logic from the core training routines in the `brittle_star_project`. + +## Quick Start + +The recommended way to use the logger is through the `get_logger()` singleton: + +```python +from experiment_logger import UnifiedLogger, get_logger + +# Initialize at the start of your script (e.g., in train.py) +logger = UnifiedLogger( + run_name="my_experiment_run", + config={"learning_rate": 3e-4}, + project_name="MyProject", + base_dir="runs", + use_wandb=True +) + +# In other files, retrieve the initialized singleton: +# logger = get_logger() + +# Log metrics (Scalar values, numpy scalars, or JAX types) +logger.log({"loss": 0.5, "accuracy": 0.98}, step=100) + +# Standard logging (Mirrored to disk and stdout) +logger.info("Training started") +logger.warning("Learning rate is very high") + +# Save checkpoints (Automatically synced to WandB as artifacts) +logger.save_checkpoint(params, step=5000) +``` + +## Logger Classes + +### `UnifiedLogger` + +The full suite for production training. It manages: +- **WandB**: Syncs metrics and uploads model checkpoints as artifacts. +- **TensorBoard**: Writes events for local visualization. +- **Local Disk**: Stores metrics in `metrics.yaml` and textual logs in `run.log`. + +### `SimpleLogger` + +A zero-dependency fallback that uses standard Python `print()` statements. Use this for standalone testing or minimal environments where you don't need persistent monitoring. + +```python +from experiment_logger import SimpleLogger +logger = SimpleLogger(run_name="test_run") +``` + +## API Features + +### `logger.progress_bar(iterable, **kwargs)` + +A smart wrapper around `tqdm` that automatically detects its environment. +- **Interactive Terminal**: Displays a normal progress bar. +- **Non-Interactive (HPC)**: Automatically disables the bar to prevent log file bloat in `slurm.out`. + +### `logger.log_non_interactive(msg: str)` + +Prints a message *only* when running in non-interactive environments. Useful for high-level progress tracking (e.g., "Epoch 5 Complete") without interactive noise. + +### `logger.save_checkpoint(params, step, prefix="checkpoint")` + +Saves model parameters using Flax serialization. +- **Local Location**: `runs//checkpoints/` +- **WandB Logic**: Automatically uploads the `.flax` file as a model artifact for lineage tracking. diff --git a/src/experiment_logger/__init__.py b/src/experiment_logger/__init__.py new file mode 100644 index 0000000..e1b2d09 --- /dev/null +++ b/src/experiment_logger/__init__.py @@ -0,0 +1,19 @@ +"""Unified logging framework for machine learning experiments. + +This package provides a unified interface for logging to multiple backends +(WandB, disk, stdout) simultaneously, ensuring no data loss. +""" + +from experiment_logger.unified_logger import UnifiedLogger, get_logger, init_logger +from experiment_logger.simple_logger import SimpleLogger +from experiment_logger.wandb_utils import finish_wandb, init_wandb + +__all__ = [ + "UnifiedLogger", + "SimpleLogger", + "get_logger", + "init_logger", + "init_wandb", + "finish_wandb", +] +__version__ = "0.1.0" diff --git a/src/experiment_logger/config_logger.py b/src/experiment_logger/config_logger.py new file mode 100644 index 0000000..fd77a28 --- /dev/null +++ b/src/experiment_logger/config_logger.py @@ -0,0 +1,36 @@ +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class LoggingConfig: + track: bool = False + wandb_project_name: str = "default-project" + wandb_entity: Optional[str] = "SEL3-2026-Groep-4" + capture_video: bool = False + + # Local Saving + save_model: bool = True # Final model + save_checkpoints: bool = True # Intermediate checkpoints + checkpoint_frequency: int = 100 + + # Remote Uploading (WandB Artifacts) + upload_final_model: bool = False + upload_checkpoints: bool = False + + hf_entity: str = "" + + def __post_init__(self): + if self.upload_final_model and not (self.track and self.save_model): + raise ValueError( + "Configuration Error: 'upload_final_model' is True, but it requires " + "both 'track' and 'save_model' to also be True." + ) + if self.upload_checkpoints and not (self.track and self.save_checkpoints): + raise ValueError( + "Configuration Error: 'upload_checkpoints' is True, but it requires " + "both 'track' and 'save_checkpoints' to also be True." + ) + + # NOTE: Checkpoint evaluation settings live under the project's + # `evaluation` config group (see brittle_star_project.configs). diff --git a/src/experiment_logger/simple_logger.py b/src/experiment_logger/simple_logger.py new file mode 100644 index 0000000..71844dc --- /dev/null +++ b/src/experiment_logger/simple_logger.py @@ -0,0 +1,85 @@ +"""Simple terminal logger for running without external backends. + +This is used for standalone package usage where WandB or TensorBoard are not desired. +It preserves the same API as UnifiedLogger but simply prints to stdout. +""" + +import logging +from typing import Any, Dict, Optional + + +class SimpleLogger: + """Simple logger that implements the UnifiedLogger interface via print statements.""" + + def __init__( + self, + run_name: str = "simple_run", + full_config: Optional[Dict[str, Any]] = None, + logging_cfg: Optional[Any] = None, + base_dir: str = "runs", + save_code: bool = False, + log_level: int = logging.INFO, + _set_as_global: bool = False, + ): + self.is_interactive = True + self.run_name = run_name + self.full_config = full_config or {} + print(f"[INIT] SimpleLogger initialized for run: {run_name}") + + def set_level(self, level: int): + pass + + def log_non_interactive(self, msg: str, *args, **kwargs): + """In SimpleLogger, we just print everything as we assume interactive use.""" + self.info(msg, *args, **kwargs) + + def progress_bar(self, iterable=None, *args, **kwargs): + """Standard tqdm wrapper that falls back to range if tqdm is missing.""" + try: + import tqdm + + return tqdm.tqdm(iterable, *args, **kwargs) + except ImportError: + return iterable + + def info(self, msg: str, *args, **kwargs): + print(f"[INFO] {msg}") + + def warning(self, msg: str, *args, **kwargs): + print(f"[WARNING] {msg}") + + def error(self, msg: str, *args, **kwargs): + print(f"[ERROR] {msg}") + + def debug(self, msg: str, *args, **kwargs): + print(f"[DEBUG] {msg}") + + def log(self, metrics: Dict[str, Any], step: Optional[int] = None, commit: bool = True): + step_str = f"Step {step}" if step is not None else "Log" + metric_str = ", ".join(f"{k}: {v}" for k, v in metrics.items()) + print(f"[{step_str}] {metric_str}") + + def save_checkpoint( + self, + params: Any, + step: int, + prefix: str = "checkpoint", + metadata: Optional[Dict[str, Any]] = None, + ): + print(f"[SAVE] Checkpoint '{prefix}' would be saved at step {step} (SimpleLogger: No-Op)") + + def save_final_model(self, params: Any, metadata: Optional[Dict[str, Any]] = None): + print("[SAVE] Final model would be saved (SimpleLogger: No-Op)") + + def sync_file(self, path: Any): + """No-op for SimpleLogger.""" + pass + + def finish(self): + print(f"[FINISH] SimpleLogger finished for run: {self.run_name}") + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.finish() diff --git a/src/experiment_logger/unified_logger.py b/src/experiment_logger/unified_logger.py new file mode 100644 index 0000000..e308d12 --- /dev/null +++ b/src/experiment_logger/unified_logger.py @@ -0,0 +1,470 @@ +"""Unified logger that writes to multiple backends simultaneously. + +This logger ensures all experimental data is preserved by writing to: +1. Weights & Biases (when available) +2. Local disk (JSON files, model checkpoints, run.log) +3. stdout (for real-time monitoring) +""" + +from enum import Enum +import logging +import yaml +import sys +import time +from pathlib import Path +from typing import Any, Dict, List, Optional + +import flax +import jax.numpy as jnp +import numpy as np + +from experiment_logger.wandb_utils import finish_wandb, init_wandb +from experiment_logger.config_logger import LoggingConfig + +# Global storage for the active logger and the proxy singleton +_active_logger: Optional[Any] = None +_proxy_instance: Optional["LoggerProxy"] = None + + +def _sanitize_for_yaml(obj: Any) -> Any: + """Convert non-primitive values into YAML-safe structures. + + In particular, avoids PyYAML serializing Enums as + ``!!python/object/apply:...`` which OmegaConf will not load. + """ + + if isinstance(obj, Enum): + return obj.name + if isinstance(obj, Path): + return str(obj) + if isinstance(obj, (np.generic, jnp.ndarray)): + try: + return obj.item() + except Exception: + pass + if isinstance(obj, np.ndarray): + return obj.tolist() + if isinstance(obj, dict): + return {str(k): _sanitize_for_yaml(v) for k, v in obj.items()} + if isinstance(obj, list): + return [_sanitize_for_yaml(v) for v in obj] + if isinstance(obj, tuple): + return [_sanitize_for_yaml(v) for v in obj] + return obj + + +def get_logger() -> "LoggerProxy": + """Retrieve the global LoggerProxy. + + This should be used for all logging calls. It returns a proxy that + delegates to the active logger (defaulting to a SimpleLogger until + init_logger is called). + """ + global _proxy_instance, _active_logger + if _proxy_instance is None: + if _active_logger is None: + # Fallback to SimpleLogger to avoid premature directory creation + from experiment_logger.simple_logger import SimpleLogger + + _active_logger = SimpleLogger(run_name="pre_init") + + _proxy_instance = LoggerProxy() + + return _proxy_instance + + +def init_logger(**kwargs) -> "UnifiedLogger": + """Initialize the full UnifiedLogger and set it as the active logger. + + This should be called once the configuration is ready. It will create + the output directories and set up all logging backends. + """ + global _active_logger + logger = UnifiedLogger(**kwargs) + _active_logger = logger + return logger + + +class LoggerProxy: + """Proxy that delegates all method calls to the active logger instance. + + This allows the logger to be swapped out (e.g., from a SimpleLogger to + a UnifiedLogger) without any clients needing to update their references. + """ + + def _get_logger(self) -> Any: + global _active_logger + if _active_logger is None: + # This shouldn't normally happen since get_logger handles it + from experiment_logger.simple_logger import SimpleLogger + + _active_logger = SimpleLogger(run_name="pre_init_fallback") + return _active_logger + + def __getattr__(self, name: str) -> Any: + return getattr(self._get_logger(), name) + + def __enter__(self): + return self._get_logger().__enter__() + + def __exit__(self, exc_type, exc_val, exc_tb): + return self._get_logger().__exit__(exc_type, exc_val, exc_tb) + + +class UnifiedLogger: + """Unified logger for scientific experiments with redundant backup.""" + + def __init__( + self, + run_name: str, + full_config: Dict[str, Any], + logging_cfg: LoggingConfig, + base_dir: str = "runs", + save_code: bool = True, + log_level: int = logging.INFO, + ): + """Initialize the unified logger. + + Args: + run_name: Unique name for this run + full_config: Full configuration dictionary with hyperparameters to be saved + logging_cfg: Structured logging configuration dataclass + base_dir: Base directory for local storage + save_code: Whether to save code to WandB + """ + self.run_name = run_name + self.full_config = full_config + self.use_wandb = logging_cfg.track + self.upload_final_model = logging_cfg.upload_final_model + self.upload_checkpoints = logging_cfg.upload_checkpoints + self.wandb_available = False + self.wandb_run = None + self.is_interactive = sys.stdout.isatty() + + # Setup local storage + self.run_dir = Path(base_dir) / run_name + self.run_dir.mkdir(parents=True, exist_ok=True) + + self.checkpoints_dir = self.run_dir / "checkpoints" + self.checkpoints_dir.mkdir(exist_ok=True) + + self.metrics_dir = self.run_dir / "metrics" + self.metrics_dir.mkdir(exist_ok=True) + + self.config_file = self.run_dir / "config.yaml" + + # Setup standard Python logging mirror + self.text_log_file = self.run_dir / "run.log" + self._text_logger = logging.getLogger(f"UnifiedLogger_{self.run_name}") + self._text_logger.setLevel(log_level) + self._text_logger.propagate = False + + # Avoid duplicate handlers if re-instantiated + if not self._text_logger.handlers: + fh = logging.FileHandler(self.text_log_file) + ch = logging.StreamHandler() + + formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s") + fh.setFormatter(formatter) + ch.setFormatter(formatter) + + self._text_logger.addHandler(fh) + self._text_logger.addHandler(ch) + + # Save config to disk + self._save_config() + + # Setup TensorBoard + self.writer = None + try: + from torch.utils.tensorboard import SummaryWriter + + self.writer = SummaryWriter(self.run_dir) + self.info("TensorBoard SummaryWriter initialized.") + except ImportError: + self.warning("tensorboard not installed. Skipping SummaryWriter.") + + # Initialize WandB if requested + if self.use_wandb: + self._init_wandb(logging_cfg.wandb_project_name, logging_cfg.wandb_entity, save_code) + + # Initialize metrics storage + self.metrics_buffer: List[Dict[str, Any]] = [] + self.step_counter = 0 + + self.info(f"Initialized UnifiedLogger for run: {run_name}") + self.info(f"Local storage: {self.run_dir.absolute()}") + self.info(f"WandB logging: {self.wandb_available}") + + def set_level(self, level: int): + """Dynamically update the verbosity of the stdout/text logger.""" + self._text_logger.setLevel(level) + + def log_non_interactive(self, msg: str, *args, **kwargs): + """Log an info message only if running in a non-interactive environment.""" + if not self.is_interactive: + self.info(msg, *args, **kwargs) + + def progress_bar(self, iterable=None, *args, **kwargs): + """Wrapper around tqdm that automatically disables in non-interactive environments.""" + import tqdm + + kwargs.setdefault("disable", not self.is_interactive) + return tqdm.tqdm(iterable, *args, **kwargs) + + def info(self, msg: str, *args, **kwargs): + """Log an info message to stdout and disk.""" + self._text_logger.info(msg, *args, **kwargs) + + def warning(self, msg: str, *args, **kwargs): + """Log a warning message to stdout and disk.""" + self._text_logger.warning(msg, *args, **kwargs) + + def error(self, msg: str, *args, **kwargs): + """Log an error message to stdout and disk.""" + self._text_logger.error(msg, *args, **kwargs) + + def debug(self, msg: str, *args, **kwargs): + """Log a debug message to stdout and disk.""" + self._text_logger.debug(msg, *args, **kwargs) + + def _init_wandb(self, project_name: str, entity: Optional[str], save_code: bool): + """Initialize Weights & Biases logging.""" + self.wandb_run = init_wandb( + project=project_name, + entity=entity, + name=self.run_name, + config=self.full_config, + save_code=save_code, + resume="allow", + ) + self.wandb_available = self.wandb_run is not None + + def _save_config(self): + """Save configuration to disk.""" + try: + with open(self.config_file, "w") as f: + yaml.safe_dump( + _sanitize_for_yaml(self.full_config), + f, + default_flow_style=False, + indent=2, + sort_keys=False, + ) + self.info(f"Config saved to {self.config_file}") + except Exception as e: + self.error(f"Error saving config: {e}") + + def log(self, metrics: Dict[str, Any], step: Optional[int] = None, commit: bool = True): + """Log metrics to all backends. + + Args: + metrics: Dictionary of metric name -> value + step: Global step counter (auto-incremented if None) + commit: Whether to commit to WandB immediately + """ + if step is None: + step = self.step_counter + self.step_counter += 1 + + # Add timestamp + metrics_with_metadata = { + "step": step, + "timestamp": time.time(), + **metrics, + } + + # Log to stdout + self._log_to_stdout(metrics_with_metadata) + + # Log to WandB + if self.wandb_run is not None: + try: + self.wandb_run.log(metrics, step=step, commit=commit) + except Exception as e: + self.warning(f"WandB logging failed: {e}") + + # Log to TensorBoard + if self.writer is not None: + for k, v in metrics.items(): + if isinstance(v, (int, float, np.floating, np.integer)): + self.writer.add_scalar(k, v, step) + elif hasattr(v, "item"): + self.writer.add_scalar(k, v.item(), step) + elif isinstance(v, (np.ndarray, jnp.ndarray)) and v.size == 1: + self.writer.add_scalar(k, v.item(), step) + + # Buffer for disk storage + self.metrics_buffer.append(metrics_with_metadata) + + # Periodically flush to disk + if len(self.metrics_buffer) >= 100: + self._flush_metrics() + + def _log_to_stdout(self, metrics: Dict[str, Any]): + """Log metrics to stdout for real-time monitoring.""" + step = metrics.get("step", "?") + metric_str = ", ".join( + f"{k}={v:.6f}" if isinstance(v, (float, np.floating)) else f"{k}={v}" + for k, v in metrics.items() + if k not in ["step", "timestamp"] + ) + self.info(f"[Step {step}] {metric_str}") + + def _flush_metrics(self): + """Flush buffered metrics to disk.""" + if not self.metrics_buffer: + return + + try: + metrics_file = self.metrics_dir / "metrics.yaml" + with open(metrics_file, "a") as f: + for metric in self.metrics_buffer: + # Convert numpy/jax types to native Python types for YAML serialization + serializable_metric = {} + for k, v in metric.items(): + if hasattr(v, "item"): # numpy/jax scalar + serializable_metric[k] = v.item() + elif isinstance(v, (np.ndarray, jnp.ndarray)): + serializable_metric[k] = v.tolist() + else: + serializable_metric[k] = v + f.write("---\n") + yaml.safe_dump( + _sanitize_for_yaml(serializable_metric), + f, + default_flow_style=False, + sort_keys=False, + ) + self.metrics_buffer.clear() + except Exception as e: + self.error(f"Error flushing metrics: {e}") + + def save_checkpoint( + self, + params: Any, + step: int, + prefix: str = "checkpoint", + metadata: Optional[Dict[str, Any]] = None, + ): + """Save model checkpoint to disk and optionally to WandB.""" + checkpoint_name = f"{prefix}_step_{step}.flax" + checkpoint_path = self.checkpoints_dir / checkpoint_name + + try: + # Save to disk using Flax serialization + with open(checkpoint_path, "wb") as f: + f.write(flax.serialization.to_bytes(params)) + + # Save metadata if provided + if metadata: + metadata_path = self.checkpoints_dir / f"{prefix}_step_{step}_metadata.yaml" + with open(metadata_path, "w") as f: + yaml.safe_dump( + _sanitize_for_yaml(metadata), + f, + default_flow_style=False, + indent=2, + sort_keys=False, + ) + + self.info(f"Checkpoint saved: {checkpoint_path}") + + # Log to WandB as artifact + if self.wandb_run is not None and self.upload_checkpoints: + try: + import wandb + + artifact = wandb.Artifact( + name=f"{self.run_name}_{prefix}", + type="model", + metadata=metadata or {}, + ) + artifact.add_file(str(checkpoint_path)) + if metadata: + artifact.add_file(str(metadata_path)) + self.wandb_run.log_artifact(artifact) + self.info("Checkpoint uploaded to WandB") + except Exception as e: + self.warning(f"Could not upload checkpoint to WandB: {e}") + + except Exception as e: + self.error(f"Error saving checkpoint: {e}") + + def save_final_model(self, params: Any, metadata: Optional[Dict[str, Any]] = None): + """Save the final trained model.""" + final_model_path = self.run_dir / "final_model.flax" + + try: + with open(final_model_path, "wb") as f: + f.write(flax.serialization.to_bytes(params)) + + if metadata: + metadata_path = self.run_dir / "final_model_metadata.yaml" + with open(metadata_path, "w") as f: + yaml.safe_dump( + _sanitize_for_yaml(metadata), + f, + default_flow_style=False, + indent=2, + sort_keys=False, + ) + + self.info(f"Final model saved: {final_model_path}") + + # Log to WandB + if self.wandb_run is not None and self.upload_final_model: + try: + import wandb + + artifact = wandb.Artifact( + name=f"{self.run_name}_final_model", + type="model", + metadata=metadata or {}, + ) + artifact.add_file(str(final_model_path)) + if metadata: + artifact.add_file(str(metadata_path)) + self.wandb_run.log_artifact(artifact) + except Exception as e: + self.warning(f"Could not upload final model to WandB: {e}") + + except Exception as e: + self.error(f"Error saving final model: {e}") + + def sync_file(self, path: Path) -> None: + """Upload a file to W&B if tracking is enabled. + + Best-effort: logs a warning on failure, never raises. + """ + if self.wandb_run is None: + return + try: + import wandb + + # "Simple sync" behavior: wandb will copy this file into the run. + wandb.save(str(path), base_path=str(path.parent)) + except Exception as e: + self.warning(f"Failed to sync file to W&B: {e}") + + def finish(self): + """Finalize logging and cleanup.""" + # Flush remaining metrics + self._flush_metrics() + + if self.writer is not None: + self.writer.close() + + self.info(f"Run complete. Results saved to: {self.run_dir.absolute()}") + + # Finish WandB run + if self.wandb_available: + finish_wandb() + + def __enter__(self): + """Context manager entry.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Context manager exit.""" + self.finish() diff --git a/src/experiment_logger/wandb_utils.py b/src/experiment_logger/wandb_utils.py new file mode 100644 index 0000000..302d308 --- /dev/null +++ b/src/experiment_logger/wandb_utils.py @@ -0,0 +1,91 @@ +"""Centralized WandB initialization utilities.""" + +import logging +from typing import Any, Dict, Optional + +logger = logging.getLogger(__name__) + + +def init_wandb( + project: str, + config: Dict[str, Any], + name: Optional[str] = None, + entity: Optional[str] = None, + sync_tensorboard: bool = False, + save_code: bool = True, + resume: str = "allow", + **kwargs, +): + """Initialize WandB with standardized settings. + + This function provides a centralized way to initialize WandB across different + scripts, ensuring consistent configuration and error handling. + + Args: + project: WandB project name + config: Configuration dictionary to log + name: Run name (auto-generated if None) + entity: WandB entity (team/user name) + sync_tensorboard: Whether to sync tensorboard logs + save_code: Whether to save code snapshots + resume: Resume strategy ("allow", "must", "never", "auto") + **kwargs: Additional arguments to pass to wandb.init() + + Returns: + wandb.Run object if successful, None otherwise + """ + try: + import wandb + import os + import sys + + # Robust HPC checking: check for API key + has_key = os.environ.get("WANDB_API_KEY") is not None + if not has_key: + try: + # Check if logged in locally via settings/netrc + has_key = wandb.setup().settings.api_key is not None + except Exception: + pass + + is_interactive = sys.stdout.isatty() + + if not has_key and not is_interactive and os.environ.get("WANDB_MODE") != "offline": + logger.warning( + "WANDB_API_KEY not found and environment is non-interactive. " + "Switching to offline mode." + ) + sync_path = f"runs/{name}" if name else "runs" + logger.warning(f"WandB is offline. Use 'wandb sync {sync_path}' to upload logs later.") + os.environ["WANDB_MODE"] = "offline" + + run = wandb.init( + project=project, + entity=entity, + name=name, + config=config, + sync_tensorboard=sync_tensorboard, + save_code=save_code, + resume=resume, + **kwargs, + ) + logger.info(f"WandB initialized successfully for project '{project}', run '{run.name}'") + return run + except ImportError: + logger.warning("WandB not installed. Skipping WandB initialization.") + return None + except Exception as e: + logger.error(f"Failed to initialize WandB: {e}") + return None + + +def finish_wandb(): + """Safely finish the current WandB run.""" + try: + import wandb + + if wandb.run is not None: + wandb.finish() + logger.info("WandB run finished successfully") + except Exception as e: + logger.warning(f"Error finishing WandB run: {e}") diff --git a/src/main.py b/src/main.py deleted file mode 100644 index ef7c36e..0000000 --- a/src/main.py +++ /dev/null @@ -1,4 +0,0 @@ -import jax - -if __name__ == "__main__": - print(jax.devices()) diff --git a/tests/.gitkeep b/tests/.gitkeep deleted file mode 100644 index e69de29..0000000 diff --git a/tests/test_adjacency.py b/tests/test_adjacency.py new file mode 100644 index 0000000..27c8b43 --- /dev/null +++ b/tests/test_adjacency.py @@ -0,0 +1,98 @@ +import jax.numpy as jnp +import numpy as np + +from brittle_star_project.MLPs import build_adjacency +from brittle_star_project.environment.env_config import MorphMode + + +def assert_symmetric(adj): + assert jnp.all(adj == adj.T) + + +def test_centralized(): + adj = build_adjacency([4, 4, 4, 4, 4], MorphMode.CENTRALIZED) + + assert adj.shape == (1, 1) + assert adj[0, 0] == 1 + + +def test_fully_connected(): + adj = build_adjacency([4, 4, 4, 4, 4], MorphMode.FULLY_CONNECTED) + + assert adj.shape == (5, 5) + assert jnp.all(adj == 1) + + +def test_ring(): + adj = build_adjacency([4, 4, 4, 4, 4], MorphMode.RING) + + assert adj.shape == (5, 5) + assert_symmetric(adj) + + # each node should connect to itself + 2 neighbors + for node in range(5): + assert adj[node, node] == 1 + assert jnp.sum(adj[node]) == 3 + neighbor1 = (node - 1) % 5 + neighbor2 = (node + 1) % 5 + assert adj[neighbor1, node] == 1 + assert adj[node, neighbor2] == 1 # Symmetrical + + +def test_segment_structure(): + segments = [4, 4, 4, 4, 4] + adj = build_adjacency(segments, MorphMode.SEGMENT) + + num_arms = 5 + num_segments = sum(segments) + num_nodes = num_arms + num_segments + + assert adj.shape == (num_nodes, num_nodes) + + # --- ring connectivity --- + for i in range(num_arms): + assert adj[i, i] == 1 + assert adj[i, (i - 1) % num_arms] == 1 + assert adj[i, (i + 1) % num_arms] == 1 + + # --- segment chain checks --- + offset = num_arms + for arm in range(5): + for i in range(4): + node = offset + arm * 4 + i + + # self + assert adj[node, node] == 1 + + # chain neighbors + if i > 0: + assert adj[node, node - 1] == 1 + if i < 3: + assert adj[node, node + 1] == 1 + + # --- ring ↔ segment connections --- + for arm in range(5): + first_seg = num_arms + arm * 4 + assert adj[arm, first_seg] == 1 + assert adj[first_seg, arm] == 1 + + save_adj(adj) + + +def save_adj(adj, name="adjacency_debug.txt"): + a = np.array(adj) + + with open(name, "w") as f: + f.write("\nAdjacency matrix:\n") + f.write(" " + " ".join([f"{i:2d}" for i in range(a.shape[0])]) + "\n") + + for i, row in enumerate(a): + line = f"{i:2d} " + " ".join(["█" if x > 0 else "." for x in row]) + f.write(line + "\n") + + +def test_no_isolated_nodes(): + adj = build_adjacency([4, 4, 4, 4, 4], MorphMode.SEGMENT) + + # no node should be completely isolated + assert jnp.all(jnp.sum(adj, axis=0) > 0) diff --git a/tests/test_configs.py b/tests/test_configs.py new file mode 100644 index 0000000..0a6a651 --- /dev/null +++ b/tests/test_configs.py @@ -0,0 +1,48 @@ +from pathlib import Path +from hydra import compose, initialize_config_dir +from omegaconf import OmegaConf +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.configs.register_configs import register_configs + +# Registration must happen before composition to enable validation against schemas +register_configs() + + +def test_config_composition_centralized(): + """Test that the centralized configuration composes and validates correctly.""" + config_dir = str(Path(__file__).parent.parent / "configs") + with initialize_config_dir(version_base="1.3", config_dir=config_dir): + # We compose the config; it follows main_config.yaml + cfg = compose(config_name="main_config", overrides=["architecture=centralized"]) + + # Merge with the structured schema and convert to a real dataclass instance + structured_cfg = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), cfg) + ) + + # Basic assertions + assert structured_cfg.architecture.name == "centralized" + assert structured_cfg.architecture.propagator is None + assert isinstance(structured_cfg.ppo.learning_rate, float) + assert structured_cfg.ppo.learning_rate > 0 + + +def test_config_composition_decentralized(): + """Test that the decentralized configuration composes and validates correctly.""" + config_dir = str(Path(__file__).parent.parent / "configs") + with initialize_config_dir(version_base="1.3", config_dir=config_dir): + cfg = compose(config_name="main_config", overrides=["architecture=decentralized"]) + + # Merge and convert to dataclass instance + structured_cfg = OmegaConf.to_object( + OmegaConf.merge(OmegaConf.structured(BrittleStarConfig), cfg) + ) + + # Basic assertions + assert structured_cfg.architecture.name == "decentralized" + assert isinstance(structured_cfg.ppo.learning_rate, float) + assert structured_cfg.ppo.learning_rate > 0 + + # Decentralized specifics + assert hasattr(structured_cfg.architecture, "message_passing_steps") + assert structured_cfg.architecture.message_passing_steps > 0 diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py new file mode 100644 index 0000000..d8be381 --- /dev/null +++ b/tests/test_evaluation.py @@ -0,0 +1,218 @@ +import numpy as np +import pytest +import yaml +from pathlib import Path + +from brittle_star_project.evaluation.checkpoint import ( + metadata_to_configs, + TrainingConfig, + load_metadata, +) +from brittle_star_project.evaluation.rollout import _maybe_clip_action +from brittle_star_project.environment.env_config import ( + MorphologyConfig, + ArenaConfig, + EnvConfig, + ObservationBoundsConfig, + MorphMode, +) +from brittle_star_project.environment.env_types import Task + + +def test_metadata_to_configs(): + """Test that a raw metadata dictionary correctly instantiates the typed configs.""" + mock_metadata = { + "morphology": { + "segments_per_arm": [4, 0, 4, 0, 0], + "use_p_control": False, + }, + "arena": {"sand_ground_color": False, "size": [15.0, 10.0]}, + "environment": { + "task": "LIGHT_ESCAPE", + "simulation_time": 5000.0, + }, + "obs_bounds": {"joint_velocity": [-10.0, 10.0]}, + } + + config = metadata_to_configs(mock_metadata) + + assert isinstance(config, TrainingConfig) + + # Check MorphologyConfig + assert isinstance(config.morphology, MorphologyConfig) + assert config.morphology.segments_per_arm == [4, 0, 4, 0, 0] + assert config.morphology.use_p_control is False + assert config.morphology.use_torque_control is False # default + + # Check ArenaConfig + assert isinstance(config.arena, ArenaConfig) + assert config.arena.sand_ground_color is False + assert config.arena.size == [15.0, 10.0] + assert config.arena.wall_height == 1.5 # default + + # Check EnvConfig + assert isinstance(config.environment, EnvConfig) + assert config.environment.task == Task.LIGHT_ESCAPE + assert config.environment.simulation_time == 5000.0 + assert config.environment.time_scale == 2 # default + + # Check ObservationBoundsConfig + assert isinstance(config.obs_bounds, ObservationBoundsConfig) + assert config.obs_bounds.joint_velocity == [-10.0, 10.0] + assert config.obs_bounds.segment_contact == [0.0, 1.0] # default + + +def test_maybe_clip_action(): + """Test action clipping against boundaries.""" + # Test valid clipping + action = np.array([1.5, -2.5, 0.0]) + low = np.array([-1.0, -1.0, -1.0]) + high = np.array([1.0, 1.0, 1.0]) + + clipped = _maybe_clip_action(action, low, high) + np.testing.assert_array_equal(clipped, np.array([1.0, -1.0, 0.0])) + + # Test skipping when bounds are None + unclipped_1 = _maybe_clip_action(action, None, high) + np.testing.assert_array_equal(unclipped_1, action) + + unclipped_2 = _maybe_clip_action(action, low, None) + np.testing.assert_array_equal(unclipped_2, action) + + # Test skipping on shape mismatch + wrong_low = np.array([-1.0, -1.0]) # Shape mismatch + unclipped_3 = _maybe_clip_action(action, wrong_low, high) + np.testing.assert_array_equal(unclipped_3, action) + + +def test_load_metadata_with_override(tmp_path: Path): + """Test that metadata can be loaded from both default and override paths.""" + # 1. Setup + model_path = tmp_path / "model.flax" + model_path.write_bytes(b"dummy") + + default_metadata_path = tmp_path / "model_metadata.yaml" + default_content = {"version": "default", "seed": 42} + with open(default_metadata_path, "w") as f: + yaml.dump(default_content, f) + + override_path = tmp_path / "custom_metadata.yaml" + override_content = {"version": "override", "seed": 1337} + with open(override_path, "w") as f: + yaml.dump(override_content, f) + + # 2. Test default behavior + loaded_default = load_metadata(model_path) + assert loaded_default == default_content + + # 3. Test override behavior + loaded_override = load_metadata(model_path, metadata_override_path=override_path) + assert loaded_override == override_content + + # 4. Test Error Case + non_existent = tmp_path / "missing.yaml" + with pytest.raises(FileNotFoundError, match="Could not find metadata YAML at"): + load_metadata(model_path, metadata_override_path=non_existent) + + +@pytest.fixture +def mock_training_config(): + return TrainingConfig( + morphology=MorphologyConfig( + segments_per_arm=[1, 1, 1, 1, 1], morph_mode=MorphMode.CENTRALIZED + ), + arena=ArenaConfig(), + environment=EnvConfig(), + obs_bounds=ObservationBoundsConfig(), + ) + + +@pytest.fixture +def mock_metadata(): + return {"architecture": {"message_passing_steps": 2}} + + +def test_build_eval_env_training_morphology(tmp_path, mock_training_config, mock_metadata): + from brittle_star_project.evaluation.eval_env_builder import build_eval_env + from unittest.mock import patch + + model_path = tmp_path / "model.flax" + patch_target = "brittle_star_project.evaluation.eval_env_builder.PolicyAgent.from_checkpoint" + with patch(patch_target) as mock_agent: + mock_agent.return_value = "mock_policy" + bundle = build_eval_env( + model_path=model_path, + training=mock_training_config, + metadata=mock_metadata, + morphology_override_path=None, + ) + assert bundle.segments_per_arm == [1, 1, 1, 1, 1] + assert bundle.num_active_arms == 5 + assert bundle.architecture == "CENTRALIZED" + assert bundle.policy == "mock_policy" + + +def test_build_eval_env_override_morphology(tmp_path, mock_training_config, mock_metadata): + from brittle_star_project.evaluation.eval_env_builder import build_eval_env + from unittest.mock import patch + + model_path = tmp_path / "model.flax" + override_path = tmp_path / "override.yaml" + override_path.write_text(yaml.dump({"segments_per_arm": [1, 0, 1, 0, 1]})) + + with patch("brittle_star_project.evaluation.eval_env_builder.PolicyAgent.from_checkpoint"): + bundle = build_eval_env( + model_path=model_path, + training=mock_training_config, + metadata=mock_metadata, + morphology_override_path=override_path, + ) + assert bundle.segments_per_arm == [1, 0, 1, 0, 1] + assert bundle.num_active_arms == 3 + # Should be smaller than 5*N + assert sum(bundle.action_mask) < len(bundle.action_mask) + + +def test_build_eval_env_action_mask_shape(tmp_path, mock_training_config, mock_metadata): + from brittle_star_project.evaluation.eval_env_builder import build_eval_env + from unittest.mock import patch + + model_path = tmp_path / "model.flax" + override_path = tmp_path / "override.yaml" + override_path.write_text(yaml.dump({"segments_per_arm": [1, 0, 1, 0, 0]})) + + with patch("brittle_star_project.evaluation.eval_env_builder.PolicyAgent.from_checkpoint"): + bundle = build_eval_env( + model_path=model_path, + training=mock_training_config, + metadata=mock_metadata, + morphology_override_path=override_path, + ) + # For each segment with P-control, there's 2 actions (pitch and yaw). + # Total segments = 5 -> 10 actions for training. + assert len(bundle.action_mask) == 10 + # Active segments = 2 -> 4 actions active. + assert sum(bundle.action_mask) == 4 + + +def test_build_eval_env_morph_mode_inherited(tmp_path, mock_training_config, mock_metadata): + from brittle_star_project.evaluation.eval_env_builder import build_eval_env + from brittle_star_project.environment.env_config import MorphMode + from unittest.mock import patch + + model_path = tmp_path / "model.flax" + override_path = tmp_path / "override.yaml" + # No morph_mode in the override YAML + override_path.write_text(yaml.dump({"segments_per_arm": [1, 0, 1, 0, 1]})) + + # Change training config to be RING + mock_training_config.morphology.morph_mode = MorphMode.RING + + with patch("brittle_star_project.evaluation.eval_env_builder.PolicyAgent.from_checkpoint"): + bundle = build_eval_env( + model_path=model_path, + training=mock_training_config, + metadata=mock_metadata, + morphology_override_path=override_path, + ) + assert bundle.architecture == "RING" diff --git a/tests/test_jax_init.py b/tests/test_jax_init.py new file mode 100644 index 0000000..e672b55 --- /dev/null +++ b/tests/test_jax_init.py @@ -0,0 +1,7 @@ +import jax + + +def test_jax_initializes(): + """Verify that JAX initializes and exposes at least one device.""" + devices = jax.devices() + assert len(devices) > 0, "JAX should expose at least one device" diff --git a/tests/test_morphology_render.py b/tests/test_morphology_render.py new file mode 100644 index 0000000..a80b5e6 --- /dev/null +++ b/tests/test_morphology_render.py @@ -0,0 +1,65 @@ +import pytest +import os +import sys + +# CRITICAL for headless cross-platform testing (devcontainers etc) +if sys.platform == "linux" and "DISPLAY" not in os.environ and "WAYLAND_DISPLAY" not in os.environ: + os.environ.setdefault("MUJOCO_GL", "egl") + +import mujoco +from PIL import Image +from brittle_star_project.environment.env_config import EnvConfig, MorphologyConfig, ArenaConfig +from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper + + +@pytest.mark.skipif(os.getenv("CI") == "true", reason="No OpenGL display in CI") +def test_render_morphologies(): + base_dir = "runs/morphologies" + os.makedirs(base_dir, exist_ok=True) + + # --- 1. Full 5-Arm Morphology --- + morph_full = MorphologyConfig(segments_per_arm=[4, 4, 4, 4, 4]) + env_full = BrittleStarJaxEnvWrapper( + morphology=morph_full, arena=ArenaConfig(), env_config=EnvConfig(), num_envs=1 + ) + state_full = env_full.reset(seed=0) + + model_full = state_full.mj_model + data_full = state_full.mj_data + + # 1. Compute forward kinematics so geoms are correctly positioned + mujoco.mj_forward(model_full, data_full) + + # 2. Render using the environment's primary camera (camera=0) + renderer_full = mujoco.Renderer(model=model_full) + renderer_full.update_scene(data_full, camera=1) + pixels_full = renderer_full.render() + image_path = os.path.join(base_dir, "5_arm.png") + Image.fromarray(pixels_full).save(image_path) + print(f"Generated full morphology render: {image_path}") + + # --- 2. Partially Amputated Morphology --- + morph_amp = MorphologyConfig(segments_per_arm=[4, 0, 4, 2, 4]) + env_amp = BrittleStarJaxEnvWrapper( + morphology=morph_amp, arena=ArenaConfig(), env_config=EnvConfig(), num_envs=1 + ) + state_amp = env_amp.reset(seed=0) + + model_amp = state_amp.mj_model + data_amp = state_amp.mj_data + + # Compute forward kinematics + mujoco.mj_forward(model_amp, data_amp) + + renderer_amp = mujoco.Renderer(model=model_amp) + renderer_amp.update_scene(data_amp, camera=1) + pixels_amp = renderer_amp.render() + image_path = os.path.join(base_dir, "amputated_arm.png") + Image.fromarray(pixels_amp).save(image_path) + print(f"Generated amputated morphology render: {image_path}") + + print("Morphology render test successful!") + + +if __name__ == "__main__": + test_render_morphologies() diff --git a/tests/test_network_shapes.py b/tests/test_network_shapes.py new file mode 100644 index 0000000..ce189f4 --- /dev/null +++ b/tests/test_network_shapes.py @@ -0,0 +1,71 @@ +import jax +import jax.numpy as jnp +from brittle_star_project.environment.env_config import MorphMode +from brittle_star_project.environment.padded_obs_wrapper import ( + compute_padding_masks, +) +from brittle_star_project.environment.obs_processing import create_obs_processor + +# We use Actor and OneDenseLayerMLP (as the critic) based on your mlps.py +from brittle_star_project.MLPs.mlps import Actor, OneDenseLayerMLP + + +def test_centralized_forward_pass_with_padding(): + batch_size = 2 + + # 1. Simulate Amputated Observation [4, 0, 4, 2, 4] -> 14 segments total + # 14 segments * 2 = 28 joints + amputated_obs = { + "joint_position": jnp.zeros((batch_size, 28)), + "joint_velocity": jnp.zeros((batch_size, 28)), + "segment_contact": jnp.zeros((batch_size, 14)), + } + + segments_per_arm = jnp.array((4, 0, 4, 2, 4)) + num_arms = jnp.where(segments_per_arm > 0, 1, 0).sum().item() + + # 2. Process and Pad Observation + masks = compute_padding_masks(segments_per_arm=list(segments_per_arm)) + obs_processor = create_obs_processor( + bounds_dict={}, + needed_copies=1, + num_arms=num_arms, + padding_masks=masks, + morph_mode=MorphMode.CENTRALIZED, + segments_per_arm=segments_per_arm, + ) + global_state = obs_processor(amputated_obs) + + # joint_position: 5 arms × 8 joints (padded) = 40 + # joint_velocity: 5 arms × 8 joints (padded) = 40 + # segment_contact: 5 arms × 4 segs (padded) = 20 + # Total = 100 (no disk or direction keys supplied) + assert global_state.shape == (batch_size, 1, 100), ( + f"Expected global state shape (2, 1, 100), got {global_state.shape}" + ) + + actor = Actor(action_dim=40) + critic = OneDenseLayerMLP() # Acts as the centralized critic + + rng = jax.random.PRNGKey(0) + rng_a, rng_c = jax.random.split(rng) + + # Initialize Flax variables + actor_params = actor.init(rng_a, global_state) + critic_params = critic.init(rng_c, global_state) + + # 5. Forward Pass Assertions + action_mean, action_log_std = actor.apply(actor_params, global_state) + value = critic.apply(critic_params, global_state) + + assert action_mean.shape == (batch_size, 1, 40), ( + f"Actor mean shape mismatch: {action_mean.shape}" + ) + assert action_log_std.shape == (40,), f"Actor log_std shape mismatch: {action_log_std.shape}" + assert value.shape == (batch_size, 1, 1) or value.shape == (batch_size,), ( + f"Critic value shape mismatch: {value.shape}" + ) + + +if __name__ == "__main__": + test_centralized_forward_pass_with_padding() diff --git a/tests/test_obs_processor.py b/tests/test_obs_processor.py new file mode 100644 index 0000000..7008648 --- /dev/null +++ b/tests/test_obs_processor.py @@ -0,0 +1,124 @@ +import jax +import jax.numpy as jnp + +from brittle_star_project.environment.obs_processing import create_obs_processor +from brittle_star_project.environment.env_config import MorphMode, ObservationBoundsConfig + + +obs_bounds = ObservationBoundsConfig().to_bounds_dict() + +# Features per decentralized agent (one arm's data): +# disk_z_tilt → scalar → 1 feat +# joint_actuator_force → 4 segs × 2 joints → 8 feat +# joint_position → 4 segs × 2 joints → 8 feat +# joint_velocity → 4 segs × 2 joints → 8 feat +# robot_direction_to_target→ (x, y) → 2 feat +# segment_contact → 4 segs → 4 feat +# Total per agent: 1+8+8+8+2+4 = 31 + +NUM_ARMS = 5 +SEGS_PER_ARM = 4 # healthy segments per arm +JOINTS_PER_SEG = 2 # from _build_joint_indices: segs * 2 + +SEGS_HEALTHY = [4, 4, 4, 4, 4] +SEGS_DAMAGED = [4, 4, 4, 4, 0] # arm 4 fully disabled +SEGS_DAMAGED_2 = [4, 0, 4, 2, 4] # arm 3 fully disabled +AGENT_INDICES = [0, 1, 2, 3, 4] + +FEAT_PER_AGENT = 1 + 8 + 8 + 8 + 2 + 4 # = 31 + +# Centralized flattening (needed_copies=1, one copy of global features): +# disk_z_tilt → repeated once → 1 feat +# joint_actuator_force → 5 arms × 8 joints → 40 feat +# joint_position → 5 arms × 8 joints → 40 feat +# joint_velocity → 5 arms × 8 joints → 40 feat +# robot_direction_to_target→ repeated once → 2 feat +# segment_contact → 5 arms × 4 segs → 20 feat +# Total: 1+40+40+40+2+20 = 143 +FEAT_CENTRALIZED = 1 + 40 + 40 + 40 + 2 + 20 # = 143 + + +def make_obs(segs_per_arm: list[int]) -> dict: + total_segs = sum(segs_per_arm) + total_joints = JOINTS_PER_SEG * total_segs + + return { + "actuator_force": jnp.ones(total_joints), + "disk_angular_velocity": jnp.zeros(3), + "disk_linear_velocity": jnp.zeros(3), + "disk_position": jnp.zeros(3), + "disk_rotation": jnp.array([0.1, 0.1, 0.5]), # (roll, pitch, yaw) + "joint_actuator_force": jnp.full(total_joints, 1.0), + "joint_position": jnp.full(total_joints, 0.5), + "joint_velocity": jnp.full(total_joints, 2.0), + "segment_contact": jnp.ones(total_segs), + "tendon_position": jnp.zeros(0), + "tendon_velocity": jnp.zeros(0), + "unit_xy_direction_to_target": jnp.array([1.0, 0.0]), + "xy_distance_to_target": jnp.array([3.5]), + } + + +def batch_obs(obs: dict): + return jax.tree_util.tree_map(lambda x: x[None, :], obs) + + +def make_processor(morph_mode: MorphMode, needed_copies: int, segments_per_arm: list[int]): + return create_obs_processor( + bounds_dict=obs_bounds, + num_arms=NUM_ARMS, + needed_copies=needed_copies, + morph_mode=morph_mode, + segments_per_arm=segments_per_arm, + agent_indices=AGENT_INDICES, + ) + + +def test_centralized_no_damage(): + proc = make_processor(MorphMode.CENTRALIZED, 1, SEGS_HEALTHY) + obs = make_obs(SEGS_HEALTHY) + obs = batch_obs(obs) + global_state = proc(obs) + + # Centralized: 5 agents flattened into 1 → shape (1, 1, 155) + assert global_state.shape == (1, 1, FEAT_CENTRALIZED) + + +def test_centralized_damaged_1_arm(): + proc = make_processor(MorphMode.CENTRALIZED, 1, SEGS_DAMAGED) + obs = make_obs(SEGS_DAMAGED) + obs = batch_obs(obs) + global_state = proc(obs) + + # shape test + assert global_state.shape == (1, 1, FEAT_CENTRALIZED) + + +def test_centralized_damaged_2_arms(): + proc = make_processor(MorphMode.CENTRALIZED, 1, SEGS_DAMAGED_2) + obs = make_obs(SEGS_DAMAGED_2) + obs = batch_obs(obs) + global_state = proc(obs) + + # shape test + assert global_state.shape == (1, 1, FEAT_CENTRALIZED) + + +def test_decentralized_fully_connected_no_damage(): + proc = make_processor(MorphMode.FULLY_CONNECTED, NUM_ARMS, SEGS_HEALTHY) + obs = make_obs(SEGS_HEALTHY) + obs = batch_obs(obs) + global_state = proc(obs) + + # shape test + assert global_state.shape == (1, NUM_ARMS, FEAT_PER_AGENT) + + +def test_decentralized_fully_connected_damaged_1_arm(): + proc = make_processor(MorphMode.FULLY_CONNECTED, NUM_ARMS, SEGS_DAMAGED) + obs = make_obs(SEGS_DAMAGED) + obs = batch_obs(obs) + global_state = proc(obs) + + # shape test + assert global_state.shape == (1, NUM_ARMS, FEAT_PER_AGENT) diff --git a/tests/test_target_direction.py b/tests/test_target_direction.py new file mode 100644 index 0000000..8317fb2 --- /dev/null +++ b/tests/test_target_direction.py @@ -0,0 +1,90 @@ +import jax.numpy as jnp + +from brittle_star_project.configs.main_config import BrittleStarConfig +from brittle_star_project.environment.env_config import MorphMode +from brittle_star_project.environment.env_types import Backend +from brittle_star_project.environment.BrittleStarJaxEnvWrapper import BrittleStarJaxEnvWrapper +from brittle_star_project.environment.obs_processing import create_obs_processor + + +def test_raw_environment_returns_allocentric_direction(): + """ + Verifies that the raw environment returns a GLOBAL (allocentric) + direction to the target. If the robot rotates in place, + the global vector to the target should remain identical. + """ + env = BrittleStarJaxEnvWrapper.default(num_envs=1, backend=Backend.MJX) + env_state = env.reset(seed=42) + raw_obs_1 = env_state.observations["unit_xy_direction_to_target"] + + ninety_deg_z_quat = jnp.array([0.7071068, 0.0, 0.0, 0.7071068]) + new_qpos = env_state.mjx_data.qpos.at[..., 3:7].set(ninety_deg_z_quat) + new_data = env_state.mjx_data.replace(qpos=new_qpos) + rotated_env_state = env_state.replace(mjx_data=new_data) + + zero_action = jnp.zeros(env.single_action_space.shape) + if len(raw_obs_1.shape) > 1: + zero_action = jnp.expand_dims(zero_action, 0) + final_env_state = env.step(rotated_env_state, zero_action) + raw_obs_2 = final_env_state.observations["unit_xy_direction_to_target"] + + # If the vector is allocentric, it should not change when the robot spins. + assert jnp.sum(jnp.abs(raw_obs_1 - raw_obs_2)) < 1e-4, ( + f"The raw environment observation changed when the robot rotated! " + f"This means it is already egocentric. " + f"Obs 1: {raw_obs_1}, Obs 2: {raw_obs_2}" + ) + + +def test_processor_converts_to_egocentric_direction(): + """ + Verifies that the obs_processor correctly applies a 2D inverse rotation + matrix to convert the global target vector into a local (egocentric) vector. + """ + cfg = BrittleStarConfig() + env = BrittleStarJaxEnvWrapper.default(num_envs=1, backend=Backend.MJX) + + segments_per_arm = jnp.array((4, 4, 4, 4, 4)) + num_arms = jnp.where(segments_per_arm > 0, 1, 0).sum().item() + + obs_processor = create_obs_processor( + bounds_dict=cfg.obs_bounds.to_bounds_dict(), + needed_copies=1, + num_arms=num_arms, + padding_masks=env.padding_masks, + morph_mode=MorphMode.CENTRALIZED, + segments_per_arm=segments_per_arm, + ) + + env_state = env.reset(seed=42) + + # --- Scenario 1 --- + # Robot is rotated 90 degrees Left (facing global Y) + # Target is straight ahead on the global X axis [1.0, 0.0] + # Because the robot is facing Y, the target on X is to its RIGHT [0.0, -1.0] locally. + dummy_obs_1 = dict(env_state.observations) + dummy_obs_1["disk_rotation"] = jnp.array([[0.0, 0.0, jnp.pi / 2.0]]) + dummy_obs_1["unit_xy_direction_to_target"] = jnp.array([[1.0, 0.0]]) + processed_1 = obs_processor(dummy_obs_1) + + # --- Scenario 2 (used to find the array indices) --- + # We change ONLY the target vector so we can isolate it in the final array + dummy_obs_2 = dict(env_state.observations) + dummy_obs_2["disk_rotation"] = jnp.array([[0.0, 0.0, jnp.pi / 2.0]]) + dummy_obs_2["unit_xy_direction_to_target"] = jnp.array([[0.0, 1.0]]) + processed_2 = obs_processor(dummy_obs_2) + + # Find the indices of the elements that changed + diff_array = jnp.abs(processed_1[0, 0] - processed_2[0, 0]) + changed_indices = jnp.where(diff_array > 1e-4)[0] + + # (143,) + local_target = processed_1[0, 0, changed_indices] + + # (2,) + expected_local_target = jnp.array([0.0, -1.0]) + + assert jnp.sum(jnp.abs(local_target - expected_local_target)) < 1e-4, ( + f"The obs_processor did not correctly rotate the vector to egocentric. " + f"Expected {expected_local_target}, but got {local_target}." + ) diff --git a/uv.lock b/uv.lock index 164ef5f..86d6ed9 100644 --- a/uv.lock +++ b/uv.lock @@ -3,45 +3,92 @@ revision = 3 requires-python = "==3.12.*" resolution-markers = [ "sys_platform == 'linux'", - "sys_platform != 'linux'", + "sys_platform == 'win32'", + "sys_platform == 'emscripten'", + "sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32'", ] [[package]] name = "2026sel3-project" version = "0.1.0" -source = { virtual = "." } +source = { editable = "." } dependencies = [ { name = "biorobot" }, + { name = "cleanrl" }, { name = "evosax" }, + { name = "flax" }, + { name = "gymnasium" }, + { name = "hydra-core" }, { name = "ipykernel" }, - { name = "jax", extra = ["cuda13"] }, + { name = "jax" }, { name = "matplotlib" }, { name = "mediapy" }, + { name = "mujoco-warp" }, + { name = "numpy" }, + { name = "optax" }, + { name = "protobuf" }, { name = "pyopengl" }, { name = "pyopengl-accelerate" }, + { name = "pyyaml" }, + { name = "torch" }, { name = "wandb" }, + { name = "warp-lang" }, +] + +[package.optional-dependencies] +analysis = [ + { name = "tensorboard" }, +] +cuda = [ + { name = "jax", extra = ["cuda13"] }, +] +evaluation = [ + { name = "imageio" }, + { name = "imageio-ffmpeg" }, ] [package.dev-dependencies] dev = [ + { name = "pre-commit" }, + { name = "pytest" }, { name = "ruff" }, ] [package.metadata] requires-dist = [ { name = "biorobot", specifier = "==0.4.2" }, + { name = "cleanrl", specifier = ">=0.4.8" }, { name = "evosax", specifier = "==0.2.0" }, + { name = "flax", specifier = ">=0.12.2" }, + { name = "gymnasium", specifier = ">=1.2.3" }, + { name = "hydra-core", specifier = ">=1.3.2" }, + { name = "imageio", marker = "extra == 'evaluation'", specifier = ">=2.35.0" }, + { name = "imageio-ffmpeg", marker = "extra == 'evaluation'", specifier = ">=0.5.1" }, { name = "ipykernel", specifier = "==7.2.0" }, - { name = "jax", extras = ["cuda13"], specifier = "==0.9.0.1" }, + { name = "jax", specifier = "==0.9.0.1" }, + { name = "jax", extras = ["cuda13"], marker = "extra == 'cuda'", specifier = "==0.9.0.1" }, { name = "matplotlib", specifier = "==3.10.8" }, { name = "mediapy", specifier = "==1.2.6" }, + { name = "mujoco-warp" }, + { name = "numpy", specifier = ">=2.0.0" }, + { name = "optax", specifier = ">=0.2.6" }, + { name = "protobuf", specifier = ">=5.0.0" }, { name = "pyopengl", specifier = ">=3.1.10" }, { name = "pyopengl-accelerate", specifier = ">=3.1.10" }, + { name = "pyyaml", specifier = ">=6.0" }, + { name = "tensorboard", marker = "extra == 'analysis'" }, + { name = "torch", specifier = ">=2.4.0" }, { name = "wandb", specifier = "==0.24.2" }, + { name = "warp-lang" }, ] +provides-extras = ["cuda", "analysis", "evaluation"] [package.metadata.requires-dev] -dev = [{ name = "ruff", specifier = ">=0.15.2" }] +dev = [ + { name = "pre-commit", specifier = ">=4.0.0" }, + { name = "pytest", specifier = ">=8.0.0" }, + { name = "ruff", specifier = ">=0.15.2" }, +] [[package]] name = "absl-py" @@ -70,6 +117,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, ] +[[package]] +name = "antlr4-python3-runtime" +version = "4.9.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034, upload-time = "2021-11-06T17:52:23.524Z" } + [[package]] name = "appnope" version = "0.1.4" @@ -90,11 +143,11 @@ wheels = [ [[package]] name = "attrs" -version = "25.4.0" +version = "26.1.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/6b/5c/685e6633917e101e5dcb62b9dd76946cbb57c26e133bae9e0cd36033c0a9/attrs-25.4.0.tar.gz", hash = "sha256:16d5969b87f0859ef33a48b35d55ac1be6e42ae49d5e853b597db70c35c57e11", size = 934251, upload-time = "2025-10-06T13:54:44.725Z" } +sdist = { url = "https://files.pythonhosted.org/packages/9a/8e/82a0fe20a541c03148528be8cac2408564a6c9a0cc7e9171802bc1d26985/attrs-26.1.0.tar.gz", hash = "sha256:d03ceb89cb322a8fd706d4fb91940737b6642aa36998fe130a9bc96c985eff32", size = 952055, upload-time = "2026-03-19T14:22:25.026Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3a/2a/7cc015f5b9f5db42b7d48157e23356022889fc354a2813c15934b7cb5c0e/attrs-25.4.0-py3-none-any.whl", hash = "sha256:adcf7e2a1fb3b36ac48d97835bb6d8ade15b8dcce26aba8bf1d14847b57a3373", size = 67615, upload-time = "2025-10-06T13:54:43.17Z" }, + { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, ] [[package]] @@ -121,11 +174,11 @@ wheels = [ [[package]] name = "certifi" -version = "2026.1.4" +version = "2026.2.25" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e0/2d/a891ca51311197f6ad14a7ef42e2399f36cf2f9bd44752b3dc4eab60fdc5/certifi-2026.1.4.tar.gz", hash = "sha256:ac726dd470482006e014ad384921ed6438c457018f4b3d204aea4281258b2120", size = 154268, upload-time = "2026-01-04T02:42:41.825Z" } +sdist = { url = "https://files.pythonhosted.org/packages/af/2d/7bf41579a8986e348fa033a31cdd0e4121114f6bce2457e8876010b092dd/certifi-2026.2.25.tar.gz", hash = "sha256:e887ab5cee78ea814d3472169153c2d12cd43b14bd03329a39a9c6e2e80bfba7", size = 155029, upload-time = "2026-02-25T02:54:17.342Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e6/ad/3cc14f097111b4de0040c83a525973216457bbeeb63739ef1ed275c1c021/certifi-2026.1.4-py3-none-any.whl", hash = "sha256:9943707519e4add1115f44c2bc244f782c0249876bf51b6599fee1ffbedd685c", size = 152900, upload-time = "2026-01-04T02:42:40.15Z" }, + { url = "https://files.pythonhosted.org/packages/9a/3c/c17fb3ca2d9c3acff52e30b309f538586f9f5b9c9cf454f3845fc9af4881/certifi-2026.2.25-py3-none-any.whl", hash = "sha256:027692e4402ad994f1c42e52a4997a9763c646b73e4096e4d5d6db8af1d6f0fa", size = 153684, upload-time = "2026-02-25T02:54:15.766Z" }, ] [[package]] @@ -152,28 +205,37 @@ wheels = [ ] [[package]] -name = "charset-normalizer" -version = "3.4.4" +name = "cfgv" +version = "3.5.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/13/69/33ddede1939fdd074bce5434295f38fae7136463422fe4fd3e0e89b98062/charset_normalizer-3.4.4.tar.gz", hash = "sha256:94537985111c35f28720e43603b8e7b43a6ecfb2ce1d3058bbe955b73404e21a", size = 129418, upload-time = "2025-10-14T04:42:32.879Z" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/b5/721b8799b04bf9afe054a3899c6cf4e880fcf8563cc71c15610242490a0c/cfgv-3.5.0.tar.gz", hash = "sha256:d5b1034354820651caa73ede66a6294d6e95c1b00acc5e9b098e917404669132", size = 7334, upload-time = "2025-11-19T20:55:51.612Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f3/85/1637cd4af66fa687396e757dec650f28025f2a2f5a5531a3208dc0ec43f2/charset_normalizer-3.4.4-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:0a98e6759f854bd25a58a73fa88833fba3b7c491169f86ce1180c948ab3fd394", size = 208425, upload-time = "2025-10-14T04:40:53.353Z" }, - { url = "https://files.pythonhosted.org/packages/9d/6a/04130023fef2a0d9c62d0bae2649b69f7b7d8d24ea5536feef50551029df/charset_normalizer-3.4.4-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b5b290ccc2a263e8d185130284f8501e3e36c5e02750fc6b6bdeb2e9e96f1e25", size = 148162, upload-time = "2025-10-14T04:40:54.558Z" }, - { url = "https://files.pythonhosted.org/packages/78/29/62328d79aa60da22c9e0b9a66539feae06ca0f5a4171ac4f7dc285b83688/charset_normalizer-3.4.4-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:74bb723680f9f7a6234dcf67aea57e708ec1fbdf5699fb91dfd6f511b0a320ef", size = 144558, upload-time = "2025-10-14T04:40:55.677Z" }, - { url = "https://files.pythonhosted.org/packages/86/bb/b32194a4bf15b88403537c2e120b817c61cd4ecffa9b6876e941c3ee38fe/charset_normalizer-3.4.4-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:f1e34719c6ed0b92f418c7c780480b26b5d9c50349e9a9af7d76bf757530350d", size = 161497, upload-time = "2025-10-14T04:40:57.217Z" }, - { url = "https://files.pythonhosted.org/packages/19/89/a54c82b253d5b9b111dc74aca196ba5ccfcca8242d0fb64146d4d3183ff1/charset_normalizer-3.4.4-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2437418e20515acec67d86e12bf70056a33abdacb5cb1655042f6538d6b085a8", size = 159240, upload-time = "2025-10-14T04:40:58.358Z" }, - { url = "https://files.pythonhosted.org/packages/c0/10/d20b513afe03acc89ec33948320a5544d31f21b05368436d580dec4e234d/charset_normalizer-3.4.4-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:11d694519d7f29d6cd09f6ac70028dba10f92f6cdd059096db198c283794ac86", size = 153471, upload-time = "2025-10-14T04:40:59.468Z" }, - { url = "https://files.pythonhosted.org/packages/61/fa/fbf177b55bdd727010f9c0a3c49eefa1d10f960e5f09d1d887bf93c2e698/charset_normalizer-3.4.4-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:ac1c4a689edcc530fc9d9aa11f5774b9e2f33f9a0c6a57864e90908f5208d30a", size = 150864, upload-time = "2025-10-14T04:41:00.623Z" }, - { url = "https://files.pythonhosted.org/packages/05/12/9fbc6a4d39c0198adeebbde20b619790e9236557ca59fc40e0e3cebe6f40/charset_normalizer-3.4.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:21d142cc6c0ec30d2efee5068ca36c128a30b0f2c53c1c07bd78cb6bc1d3be5f", size = 150647, upload-time = "2025-10-14T04:41:01.754Z" }, - { url = "https://files.pythonhosted.org/packages/ad/1f/6a9a593d52e3e8c5d2b167daf8c6b968808efb57ef4c210acb907c365bc4/charset_normalizer-3.4.4-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:5dbe56a36425d26d6cfb40ce79c314a2e4dd6211d51d6d2191c00bed34f354cc", size = 145110, upload-time = "2025-10-14T04:41:03.231Z" }, - { url = "https://files.pythonhosted.org/packages/30/42/9a52c609e72471b0fc54386dc63c3781a387bb4fe61c20231a4ebcd58bdd/charset_normalizer-3.4.4-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:5bfbb1b9acf3334612667b61bd3002196fe2a1eb4dd74d247e0f2a4d50ec9bbf", size = 162839, upload-time = "2025-10-14T04:41:04.715Z" }, - { url = "https://files.pythonhosted.org/packages/c4/5b/c0682bbf9f11597073052628ddd38344a3d673fda35a36773f7d19344b23/charset_normalizer-3.4.4-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:d055ec1e26e441f6187acf818b73564e6e6282709e9bcb5b63f5b23068356a15", size = 150667, upload-time = "2025-10-14T04:41:05.827Z" }, - { url = "https://files.pythonhosted.org/packages/e4/24/a41afeab6f990cf2daf6cb8c67419b63b48cf518e4f56022230840c9bfb2/charset_normalizer-3.4.4-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:af2d8c67d8e573d6de5bc30cdb27e9b95e49115cd9baad5ddbd1a6207aaa82a9", size = 160535, upload-time = "2025-10-14T04:41:06.938Z" }, - { url = "https://files.pythonhosted.org/packages/2a/e5/6a4ce77ed243c4a50a1fecca6aaaab419628c818a49434be428fe24c9957/charset_normalizer-3.4.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:780236ac706e66881f3b7f2f32dfe90507a09e67d1d454c762cf642e6e1586e0", size = 154816, upload-time = "2025-10-14T04:41:08.101Z" }, - { url = "https://files.pythonhosted.org/packages/a8/ef/89297262b8092b312d29cdb2517cb1237e51db8ecef2e9af5edbe7b683b1/charset_normalizer-3.4.4-cp312-cp312-win32.whl", hash = "sha256:5833d2c39d8896e4e19b689ffc198f08ea58116bee26dea51e362ecc7cd3ed26", size = 99694, upload-time = "2025-10-14T04:41:09.23Z" }, - { url = "https://files.pythonhosted.org/packages/3d/2d/1e5ed9dd3b3803994c155cd9aacb60c82c331bad84daf75bcb9c91b3295e/charset_normalizer-3.4.4-cp312-cp312-win_amd64.whl", hash = "sha256:a79cfe37875f822425b89a82333404539ae63dbdddf97f84dcbc3d339aae9525", size = 107131, upload-time = "2025-10-14T04:41:10.467Z" }, - { url = "https://files.pythonhosted.org/packages/d0/d9/0ed4c7098a861482a7b6a95603edce4c0d9db2311af23da1fb2b75ec26fc/charset_normalizer-3.4.4-cp312-cp312-win_arm64.whl", hash = "sha256:376bec83a63b8021bb5c8ea75e21c4ccb86e7e45ca4eb81146091b56599b80c3", size = 100390, upload-time = "2025-10-14T04:41:11.915Z" }, - { url = "https://files.pythonhosted.org/packages/0a/4c/925909008ed5a988ccbb72dcc897407e5d6d3bd72410d69e051fc0c14647/charset_normalizer-3.4.4-py3-none-any.whl", hash = "sha256:7a32c560861a02ff789ad905a2fe94e3f840803362c84fecf1851cb4cf3dc37f", size = 53402, upload-time = "2025-10-14T04:42:31.76Z" }, + { url = "https://files.pythonhosted.org/packages/db/3c/33bac158f8ab7f89b2e59426d5fe2e4f63f7ed25df84c036890172b412b5/cfgv-3.5.0-py2.py3-none-any.whl", hash = "sha256:a8dc6b26ad22ff227d2634a65cb388215ce6cc96bbcc5cfde7641ae87e8dacc0", size = 7445, upload-time = "2025-11-19T20:55:50.744Z" }, +] + +[[package]] +name = "charset-normalizer" +version = "3.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7b/60/e3bec1881450851b087e301bedc3daa9377a4d45f1c26aa90b0b235e38aa/charset_normalizer-3.4.6.tar.gz", hash = "sha256:1ae6b62897110aa7c79ea2f5dd38d1abca6db663687c0b1ad9aed6f6bae3d9d6", size = 143363, upload-time = "2026-03-15T18:53:25.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e5/62/c0815c992c9545347aeea7859b50dc9044d147e2e7278329c6e02ac9a616/charset_normalizer-3.4.6-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:2ef7fedc7a6ecbe99969cd09632516738a97eeb8bd7258bf8a0f23114c057dab", size = 295154, upload-time = "2026-03-15T18:50:50.88Z" }, + { url = "https://files.pythonhosted.org/packages/a8/37/bdca6613c2e3c58c7421891d80cc3efa1d32e882f7c4a7ee6039c3fc951a/charset_normalizer-3.4.6-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a4ea868bc28109052790eb2b52a9ab33f3aa7adc02f96673526ff47419490e21", size = 199191, upload-time = "2026-03-15T18:50:52.658Z" }, + { url = "https://files.pythonhosted.org/packages/6c/92/9934d1bbd69f7f398b38c5dae1cbf9cc672e7c34a4adf7b17c0a9c17d15d/charset_normalizer-3.4.6-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:836ab36280f21fc1a03c99cd05c6b7af70d2697e374c7af0b61ed271401a72a2", size = 218674, upload-time = "2026-03-15T18:50:54.102Z" }, + { url = "https://files.pythonhosted.org/packages/af/90/25f6ab406659286be929fd89ab0e78e38aa183fc374e03aa3c12d730af8a/charset_normalizer-3.4.6-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f1ce721c8a7dfec21fcbdfe04e8f68174183cf4e8188e0645e92aa23985c57ff", size = 215259, upload-time = "2026-03-15T18:50:55.616Z" }, + { url = "https://files.pythonhosted.org/packages/4e/ef/79a463eb0fff7f96afa04c1d4c51f8fc85426f918db467854bfb6a569ce3/charset_normalizer-3.4.6-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0e28d62a8fc7a1fa411c43bd65e346f3bce9716dc51b897fbe930c5987b402d5", size = 207276, upload-time = "2026-03-15T18:50:57.054Z" }, + { url = "https://files.pythonhosted.org/packages/f7/72/d0426afec4b71dc159fa6b4e68f868cd5a3ecd918fec5813a15d292a7d10/charset_normalizer-3.4.6-cp312-cp312-manylinux_2_31_armv7l.whl", hash = "sha256:530d548084c4a9f7a16ed4a294d459b4f229db50df689bfe92027452452943a0", size = 195161, upload-time = "2026-03-15T18:50:58.686Z" }, + { url = "https://files.pythonhosted.org/packages/bf/18/c82b06a68bfcb6ce55e508225d210c7e6a4ea122bfc0748892f3dc4e8e11/charset_normalizer-3.4.6-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:30f445ae60aad5e1f8bdbb3108e39f6fbc09f4ea16c815c66578878325f8f15a", size = 203452, upload-time = "2026-03-15T18:51:00.196Z" }, + { url = "https://files.pythonhosted.org/packages/44/d6/0c25979b92f8adafdbb946160348d8d44aa60ce99afdc27df524379875cb/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ac2393c73378fea4e52aa56285a3d64be50f1a12395afef9cce47772f60334c2", size = 202272, upload-time = "2026-03-15T18:51:01.703Z" }, + { url = "https://files.pythonhosted.org/packages/2e/3d/7fea3e8fe84136bebbac715dd1221cc25c173c57a699c030ab9b8900cbb7/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:90ca27cd8da8118b18a52d5f547859cc1f8354a00cd1e8e5120df3e30d6279e5", size = 195622, upload-time = "2026-03-15T18:51:03.526Z" }, + { url = "https://files.pythonhosted.org/packages/57/8a/d6f7fd5cb96c58ef2f681424fbca01264461336d2a7fc875e4446b1f1346/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:8e5a94886bedca0f9b78fecd6afb6629142fd2605aa70a125d49f4edc6037ee6", size = 220056, upload-time = "2026-03-15T18:51:05.269Z" }, + { url = "https://files.pythonhosted.org/packages/16/50/478cdda782c8c9c3fb5da3cc72dd7f331f031e7f1363a893cdd6ca0f8de0/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:695f5c2823691a25f17bc5d5ffe79fa90972cc34b002ac6c843bb8a1720e950d", size = 203751, upload-time = "2026-03-15T18:51:06.858Z" }, + { url = "https://files.pythonhosted.org/packages/75/fc/cc2fcac943939c8e4d8791abfa139f685e5150cae9f94b60f12520feaa9b/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:231d4da14bcd9301310faf492051bee27df11f2bc7549bc0bb41fef11b82daa2", size = 216563, upload-time = "2026-03-15T18:51:08.564Z" }, + { url = "https://files.pythonhosted.org/packages/a8/b7/a4add1d9a5f68f3d037261aecca83abdb0ab15960a3591d340e829b37298/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a056d1ad2633548ca18ffa2f85c202cfb48b68615129143915b8dc72a806a923", size = 209265, upload-time = "2026-03-15T18:51:10.312Z" }, + { url = "https://files.pythonhosted.org/packages/6c/18/c094561b5d64a24277707698e54b7f67bd17a4f857bbfbb1072bba07c8bf/charset_normalizer-3.4.6-cp312-cp312-win32.whl", hash = "sha256:c2274ca724536f173122f36c98ce188fd24ce3dad886ec2b7af859518ce008a4", size = 144229, upload-time = "2026-03-15T18:51:11.694Z" }, + { url = "https://files.pythonhosted.org/packages/ab/20/0567efb3a8fd481b8f34f739ebddc098ed062a59fed41a8d193a61939e8f/charset_normalizer-3.4.6-cp312-cp312-win_amd64.whl", hash = "sha256:c8ae56368f8cc97c7e40a7ee18e1cedaf8e780cd8bc5ed5ac8b81f238614facb", size = 154277, upload-time = "2026-03-15T18:51:13.004Z" }, + { url = "https://files.pythonhosted.org/packages/15/57/28d79b44b51933119e21f65479d0864a8d5893e494cf5daab15df0247c17/charset_normalizer-3.4.6-cp312-cp312-win_arm64.whl", hash = "sha256:899d28f422116b08be5118ef350c292b36fc15ec2daeb9ea987c89281c7bb5c4", size = 142817, upload-time = "2026-03-15T18:51:14.408Z" }, + { url = "https://files.pythonhosted.org/packages/2a/68/687187c7e26cb24ccbd88e5069f5ef00eba804d36dde11d99aad0838ab45/charset_normalizer-3.4.6-py3-none-any.whl", hash = "sha256:947cf925bc916d90adba35a64c82aace04fa39b46b52d4630ece166655905a69", size = 61455, upload-time = "2026-03-15T18:53:23.833Z" }, ] [[package]] @@ -193,6 +255,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/12/0c/96102c01dd02ae740d4afc3644d5c7d7fc51d3feefd67300a2aa1ddbf7cb/chex-0.1.91-py3-none-any.whl", hash = "sha256:6fc4cbfc22301c08d4a7ef706045668410100962eba8ba6af03fa07f4e5dcf9b", size = 100965, upload-time = "2025-09-01T21:49:31.141Z" }, ] +[[package]] +name = "cleanrl" +version = "0.4.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "gym" }, + { name = "seaborn" }, + { name = "stable-baselines3" }, + { name = "tensorboard" }, + { name = "torch" }, + { name = "wandb" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/87/e6/033795b58e14de7f0c294e3bdf571d9491eca6b467c6c0bc7fdcb05040ed/cleanrl-0.4.8.tar.gz", hash = "sha256:947a5b4c006f43cf90be03b90c5345d118468260c7269c11e5d795954a2effa5", size = 89200, upload-time = "2021-05-16T02:42:40.292Z" } + [[package]] name = "click" version = "8.3.1" @@ -254,6 +330,69 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/e2/f05240d2c39a1ed228d8328a78b6f44cd695f7ef47beb3e684cf93604f86/contourpy-1.3.3-cp312-cp312-win_arm64.whl", hash = "sha256:07ce5ed73ecdc4a03ffe3e1b3e3c1166db35ae7584be76f65dbbe28a7791b0cc", size = 193655, upload-time = "2025-07-26T12:01:37.999Z" }, ] +[[package]] +name = "cuda-bindings" +version = "13.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-pathfinder", marker = "sys_platform == 'linux'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/52/c8/b2589d68acf7e3d63e2be330b84bc25712e97ed799affbca7edd7eae25d6/cuda_bindings-13.2.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e865447abfb83d6a98ad5130ed3c70b1fc295ae3eeee39fd07b4ddb0671b6788", size = 5722404, upload-time = "2026-03-11T00:12:44.041Z" }, + { url = "https://files.pythonhosted.org/packages/1f/92/f899f7bbb5617bb65ec52a6eac1e9a1447a86b916c4194f8a5001b8cde0c/cuda_bindings-13.2.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:46d8776a55d6d5da9dd6e9858fba2efcda2abe6743871dee47dd06eb8cb6d955", size = 6320619, upload-time = "2026-03-11T00:12:45.939Z" }, +] + +[[package]] +name = "cuda-pathfinder" +version = "1.5.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/93/66/0c02bd330e7d976f83fa68583d6198d76f23581bcbb5c0e98a6148f326e5/cuda_pathfinder-1.5.0-py3-none-any.whl", hash = "sha256:498f90a9e9de36044a7924742aecce11c50c49f735f1bc53e05aa46de9ea4110", size = 49739, upload-time = "2026-03-24T21:14:30.869Z" }, +] + +[[package]] +name = "cuda-toolkit" +version = "13.0.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/57/b2/453099f5f3b698d7d0eab38916aac44c7f76229f451709e2eb9db6615dcd/cuda_toolkit-13.0.2-py2.py3-none-any.whl", hash = "sha256:b198824cf2f54003f50d64ada3a0f184b42ca0846c1c94192fa269ecd97a66eb", size = 2364, upload-time = "2025-12-19T23:24:07.328Z" }, +] + +[package.optional-dependencies] +cublas = [ + { name = "nvidia-cublas", marker = "sys_platform == 'linux'" }, +] +cudart = [ + { name = "nvidia-cuda-runtime", marker = "sys_platform == 'linux'" }, +] +cufft = [ + { name = "nvidia-cufft", marker = "sys_platform == 'linux'" }, +] +cufile = [ + { name = "nvidia-cufile", marker = "sys_platform == 'linux'" }, +] +cupti = [ + { name = "nvidia-cuda-cupti", marker = "sys_platform == 'linux'" }, +] +curand = [ + { name = "nvidia-curand", marker = "sys_platform == 'linux'" }, +] +cusolver = [ + { name = "nvidia-cusolver", marker = "sys_platform == 'linux'" }, +] +cusparse = [ + { name = "nvidia-cusparse", marker = "sys_platform == 'linux'" }, +] +nvjitlink = [ + { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux'" }, +] +nvrtc = [ + { name = "nvidia-cuda-nvrtc", marker = "sys_platform == 'linux'" }, +] +nvtx = [ + { name = "nvidia-nvtx", marker = "sys_platform == 'linux'" }, +] + [[package]] name = "cycler" version = "0.12.1" @@ -263,6 +402,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, ] +[[package]] +name = "dataclasses-json" +version = "0.6.7" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "marshmallow" }, + { name = "typing-inspect" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/64/a4/f71d9cf3a5ac257c993b5ca3f93df5f7fb395c725e7f1e6479d2514173c3/dataclasses_json-0.6.7.tar.gz", hash = "sha256:b6b3e528266ea45b9535223bc53ca645f5208833c29229e847b3f26a1cc55fc0", size = 32227, upload-time = "2024-06-09T16:20:19.103Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/be/d0d44e092656fe7a06b55e6103cbce807cdbdee17884a5367c68c9860853/dataclasses_json-0.6.7-py3-none-any.whl", hash = "sha256:0dbf33f26c8d5305befd61b39d2b3414e8a407bedc2834dea9b8d642666fb40a", size = 28686, upload-time = "2024-06-09T16:20:16.715Z" }, +] + [[package]] name = "debugpy" version = "1.8.20" @@ -285,9 +437,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190, upload-time = "2025-02-24T04:41:32.565Z" }, ] +[[package]] +name = "distlib" +version = "0.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/96/8e/709914eb2b5749865801041647dc7f4e6d00b549cfe88b65ca192995f07c/distlib-0.4.0.tar.gz", hash = "sha256:feec40075be03a04501a973d81f633735b4b69f98b05450592310c0f401a4e0d", size = 614605, upload-time = "2025-07-17T16:52:00.465Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/33/6b/e0547afaf41bf2c42e52430072fa5658766e3d65bd4b03a563d1b6336f57/distlib-0.4.0-py2.py3-none-any.whl", hash = "sha256:9659f7d87e46584a30b5780e43ac7a2143098441670ff0a49d5f9034c54a6c16", size = 469047, upload-time = "2025-07-17T16:51:58.613Z" }, +] + [[package]] name = "dm-control" -version = "1.0.37" +version = "1.0.38" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "absl-py" }, @@ -306,9 +467,9 @@ dependencies = [ { name = "setuptools" }, { name = "tqdm" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/50/61/ec736b3d40134d4cf4e635b539b9967ad5c8c92b60f0da1b6eb83dffba1b/dm_control-1.0.37.tar.gz", hash = "sha256:3327e64538c230b8b95db5ed38e630236b81ae1775463eec27932aa22f730f44", size = 56273956, upload-time = "2026-02-13T10:51:24.157Z" } +sdist = { url = "https://files.pythonhosted.org/packages/7c/d7/ff2c4703c43c3cf14aa092bfb9a5eb740849eee61c25c0d5aed6a57985a1/dm_control-1.0.38.tar.gz", hash = "sha256:f5966799662d8914bb4f181d2be939da2079f0e08e198f901bb54e61d46c0b90", size = 56274133, upload-time = "2026-03-11T08:31:50.853Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/78/59/0977297d708173b0527cc3b4c4b84222c3c3bfa6a644113d2ef2163ff710/dm_control-1.0.37-py3-none-any.whl", hash = "sha256:d30573dec834201ec7fa15117df32f52012448e1b035258c69150e6dbab9b4a9", size = 56446231, upload-time = "2026-02-13T10:51:20.39Z" }, + { url = "https://files.pythonhosted.org/packages/ec/0e/2bed3a6d225921e31bf9200a9c4470f7bc3a12b940bfbca22fa7aa9dfa0e/dm_control-1.0.38-py3-none-any.whl", hash = "sha256:ddb71c4b360a503eabf3c62ba151a0d9404cbe22b23561202b90310558690371", size = 56446251, upload-time = "2026-03-11T08:31:45.102Z" }, ] [[package]] @@ -354,17 +515,16 @@ wheels = [ [[package]] name = "etils" -version = "1.13.0" +version = "1.14.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/9b/a0/522bbff0f3cdd37968f90dd7f26c7aa801ed87f5ba335f156de7f2b88a48/etils-1.13.0.tar.gz", hash = "sha256:a5b60c71f95bcd2d43d4e9fb3dc3879120c1f60472bb5ce19f7a860b1d44f607", size = 106368, upload-time = "2025-07-15T10:29:10.563Z" } +sdist = { url = "https://files.pythonhosted.org/packages/26/ce/6e067242fde898841922ac6fc82b0bb2fe35c38e995880bdffdfbe30182a/etils-1.14.0.tar.gz", hash = "sha256:8136e7f4c4173cd0af0ca5481c4475152f0b8686192951eefa60ee8711e1ede4", size = 108127, upload-time = "2026-03-04T17:41:36.291Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e7/98/87b5946356095738cb90a6df7b35ff69ac5750f6e783d5fbcc5cb3b6cbd7/etils-1.13.0-py3-none-any.whl", hash = "sha256:d9cd4f40fbe77ad6613b7348a18132cc511237b6c076dbb89105c0b520a4c6bb", size = 170603, upload-time = "2025-07-15T10:29:09.076Z" }, + { url = "https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl", hash = "sha256:b5df7341f54dbe1405a4450b2741207b4a8c279780402b45f87202b94dfc52b4", size = 172934, upload-time = "2026-03-04T17:41:35.01Z" }, ] [package.optional-dependencies] epath = [ { name = "fsspec" }, - { name = "importlib-resources" }, { name = "typing-extensions" }, { name = "zipp" }, ] @@ -406,9 +566,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/05/2c/ffc08c54c05cdce6fbed2aeebc46348dbe180c6d2c541c7af7ba0aa5f5f8/Farama_Notifications-0.0.4-py3-none-any.whl", hash = "sha256:14de931035a41961f7c056361dc7f980762a143d05791ef5794a751a2caf05ae", size = 2511, upload-time = "2023-02-27T18:28:39.447Z" }, ] +[[package]] +name = "filelock" +version = "3.25.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/b8/00651a0f559862f3bb7d6f7477b192afe3f583cc5e26403b44e59a55ab34/filelock-3.25.2.tar.gz", hash = "sha256:b64ece2b38f4ca29dd3e810287aa8c48182bbecd1ae6e9ae126c9b35f1382694", size = 40480, upload-time = "2026-03-11T20:45:38.487Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/a5/842ae8f0c08b61d6484b52f99a03510a3a72d23141942d216ebe81fefbce/filelock-3.25.2-py3-none-any.whl", hash = "sha256:ca8afb0da15f229774c9ad1b455ed96e85a81373065fb10446672f64444ddf70", size = 26759, upload-time = "2026-03-11T20:45:37.437Z" }, +] + [[package]] name = "flax" -version = "0.12.2" +version = "0.12.6" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jax" }, @@ -416,32 +585,33 @@ dependencies = [ { name = "numpy" }, { name = "optax" }, { name = "orbax-checkpoint" }, + { name = "orbax-export" }, { name = "pyyaml" }, { name = "rich" }, { name = "tensorstore" }, { name = "treescope" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/6b/7e/c4c66ab9b41149cf7a1961907d9a844832af1e76b121b35235a618c92825/flax-0.12.2.tar.gz", hash = "sha256:e9723b0881e571abe61885bb8770f53fdb3c383b6b3f5a923dcf6f1e9a687905", size = 5008370, upload-time = "2025-12-18T22:36:19.988Z" } +sdist = { url = "https://files.pythonhosted.org/packages/9a/40/d9707f22377d34dc9eaa5df67e51db4d667db9538b0f2c60c0921bc86473/flax-0.12.6.tar.gz", hash = "sha256:309a5fdfac8fe9cc03260c122a2cab6881bc366cd2d928aedb80ddffbfb202e4", size = 5077551, upload-time = "2026-03-20T21:10:22.661Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b7/6b/7b75508251f4220df8f68e7718b476ee3d614a2a51f9eace97393ee91b46/flax-0.12.2-py3-none-any.whl", hash = "sha256:912fdd8a7c623ec8b2694b28d2827608e7fc82a3a6f8fff17ec5038f2bca66f4", size = 488031, upload-time = "2025-12-18T22:36:18.01Z" }, + { url = "https://files.pythonhosted.org/packages/32/0d/aa360056c4dbb263339aa4d315c45b2c7046ef95f7b2f55732eed396a63f/flax-0.12.6-py3-none-any.whl", hash = "sha256:c16e7ea1daa96153b6cc91e1e8274fa7cdb36c80180038b7e8ddb9b4e93c80f1", size = 516706, upload-time = "2026-03-20T21:10:20.683Z" }, ] [[package]] name = "fonttools" -version = "4.61.1" +version = "4.62.1" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ec/ca/cf17b88a8df95691275a3d77dc0a5ad9907f328ae53acbe6795da1b2f5ed/fonttools-4.61.1.tar.gz", hash = "sha256:6675329885c44657f826ef01d9e4fb33b9158e9d93c537d84ad8399539bc6f69", size = 3565756, upload-time = "2025-12-12T17:31:24.246Z" } +sdist = { url = "https://files.pythonhosted.org/packages/9a/08/7012b00a9a5874311b639c3920270c36ee0c445b69d9989a85e5c92ebcb0/fonttools-4.62.1.tar.gz", hash = "sha256:e54c75fd6041f1122476776880f7c3c3295ffa31962dc6ebe2543c00dca58b5d", size = 3580737, upload-time = "2026-03-13T13:54:25.52Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/6f/16/7decaa24a1bd3a70c607b2e29f0adc6159f36a7e40eaba59846414765fd4/fonttools-4.61.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:f3cb4a569029b9f291f88aafc927dd53683757e640081ca8c412781ea144565e", size = 2851593, upload-time = "2025-12-12T17:30:04.225Z" }, - { url = "https://files.pythonhosted.org/packages/94/98/3c4cb97c64713a8cf499b3245c3bf9a2b8fd16a3e375feff2aed78f96259/fonttools-4.61.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:41a7170d042e8c0024703ed13b71893519a1a6d6e18e933e3ec7507a2c26a4b2", size = 2400231, upload-time = "2025-12-12T17:30:06.47Z" }, - { url = "https://files.pythonhosted.org/packages/b7/37/82dbef0f6342eb01f54bca073ac1498433d6ce71e50c3c3282b655733b31/fonttools-4.61.1-cp312-cp312-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:10d88e55330e092940584774ee5e8a6971b01fc2f4d3466a1d6c158230880796", size = 4954103, upload-time = "2025-12-12T17:30:08.432Z" }, - { url = "https://files.pythonhosted.org/packages/6c/44/f3aeac0fa98e7ad527f479e161aca6c3a1e47bb6996b053d45226fe37bf2/fonttools-4.61.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:15acc09befd16a0fb8a8f62bc147e1a82817542d72184acca9ce6e0aeda9fa6d", size = 5004295, upload-time = "2025-12-12T17:30:10.56Z" }, - { url = "https://files.pythonhosted.org/packages/14/e8/7424ced75473983b964d09f6747fa09f054a6d656f60e9ac9324cf40c743/fonttools-4.61.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e6bcdf33aec38d16508ce61fd81838f24c83c90a1d1b8c68982857038673d6b8", size = 4944109, upload-time = "2025-12-12T17:30:12.874Z" }, - { url = "https://files.pythonhosted.org/packages/c8/8b/6391b257fa3d0b553d73e778f953a2f0154292a7a7a085e2374b111e5410/fonttools-4.61.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:5fade934607a523614726119164ff621e8c30e8fa1ffffbbd358662056ba69f0", size = 5093598, upload-time = "2025-12-12T17:30:15.79Z" }, - { url = "https://files.pythonhosted.org/packages/d9/71/fd2ea96cdc512d92da5678a1c98c267ddd4d8c5130b76d0f7a80f9a9fde8/fonttools-4.61.1-cp312-cp312-win32.whl", hash = "sha256:75da8f28eff26defba42c52986de97b22106cb8f26515b7c22443ebc9c2d3261", size = 2269060, upload-time = "2025-12-12T17:30:18.058Z" }, - { url = "https://files.pythonhosted.org/packages/80/3b/a3e81b71aed5a688e89dfe0e2694b26b78c7d7f39a5ffd8a7d75f54a12a8/fonttools-4.61.1-cp312-cp312-win_amd64.whl", hash = "sha256:497c31ce314219888c0e2fce5ad9178ca83fe5230b01a5006726cdf3ac9f24d9", size = 2319078, upload-time = "2025-12-12T17:30:22.862Z" }, - { url = "https://files.pythonhosted.org/packages/c7/4e/ce75a57ff3aebf6fc1f4e9d508b8e5810618a33d900ad6c19eb30b290b97/fonttools-4.61.1-py3-none-any.whl", hash = "sha256:17d2bf5d541add43822bcf0c43d7d847b160c9bb01d15d5007d84e2217aaa371", size = 1148996, upload-time = "2025-12-12T17:31:21.03Z" }, + { url = "https://files.pythonhosted.org/packages/47/d4/dbacced3953544b9a93088cc10ef2b596d348c983d5c67a404fa41ec51ba/fonttools-4.62.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:90365821debbd7db678809c7491ca4acd1e0779b9624cdc6ddaf1f31992bf974", size = 2870219, upload-time = "2026-03-13T13:52:53.664Z" }, + { url = "https://files.pythonhosted.org/packages/66/9e/a769c8e99b81e5a87ab7e5e7236684de4e96246aae17274e5347d11ebd78/fonttools-4.62.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:12859ff0b47dd20f110804c3e0d0970f7b832f561630cd879969011541a464a9", size = 2414891, upload-time = "2026-03-13T13:52:56.493Z" }, + { url = "https://files.pythonhosted.org/packages/69/64/f19a9e3911968c37e1e620e14dfc5778299e1474f72f4e57c5ec771d9489/fonttools-4.62.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9c125ffa00c3d9003cdaaf7f2c79e6e535628093e14b5de1dccb08859b680936", size = 5033197, upload-time = "2026-03-13T13:52:59.179Z" }, + { url = "https://files.pythonhosted.org/packages/9b/8a/99c8b3c3888c5c474c08dbfd7c8899786de9604b727fcefb055b42c84bba/fonttools-4.62.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:149f7d84afca659d1a97e39a4778794a2f83bf344c5ee5134e09995086cc2392", size = 4988768, upload-time = "2026-03-13T13:53:02.761Z" }, + { url = "https://files.pythonhosted.org/packages/d1/c6/0f904540d3e6ab463c1243a0d803504826a11604c72dd58c2949796a1762/fonttools-4.62.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:0aa72c43a601cfa9273bb1ae0518f1acadc01ee181a6fc60cd758d7fdadffc04", size = 4971512, upload-time = "2026-03-13T13:53:05.678Z" }, + { url = "https://files.pythonhosted.org/packages/29/0b/5cbef6588dc9bd6b5c9ad6a4d5a8ca384d0cea089da31711bbeb4f9654a6/fonttools-4.62.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:19177c8d96c7c36359266e571c5173bcee9157b59cfc8cb0153c5673dc5a3a7d", size = 5122723, upload-time = "2026-03-13T13:53:08.662Z" }, + { url = "https://files.pythonhosted.org/packages/4a/47/b3a5342d381595ef439adec67848bed561ab7fdb1019fa522e82101b7d9c/fonttools-4.62.1-cp312-cp312-win32.whl", hash = "sha256:a24decd24d60744ee8b4679d38e88b8303d86772053afc29b19d23bb8207803c", size = 2281278, upload-time = "2026-03-13T13:53:10.998Z" }, + { url = "https://files.pythonhosted.org/packages/28/b1/0c2ab56a16f409c6c8a68816e6af707827ad5d629634691ff60a52879792/fonttools-4.62.1-cp312-cp312-win_amd64.whl", hash = "sha256:9e7863e10b3de72376280b515d35b14f5eeed639d1aa7824f4cf06779ec65e42", size = 2331414, upload-time = "2026-03-13T13:53:13.992Z" }, + { url = "https://files.pythonhosted.org/packages/fd/ba/56147c165442cc5ba7e82ecf301c9a68353cede498185869e6e02b4c264f/fonttools-4.62.1-py3-none-any.whl", hash = "sha256:7487782e2113861f4ddcc07c3436450659e3caa5e470b27dc2177cade2d8e7fd", size = 1152647, upload-time = "2026-03-13T13:54:22.735Z" }, ] [[package]] @@ -511,6 +681,55 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5a/3f/efeb7c6801c46e11bd666a5180f0d615f74f72264212f74f39586c6fda9d/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-manylinux_2_28_x86_64.whl", hash = "sha256:ce6724bb7cb3d0543dcba17206dce909f94176e68220b8eafee72e9f92bcf542", size = 243522, upload-time = "2026-01-28T05:58:03.517Z" }, { url = "https://files.pythonhosted.org/packages/cf/b9/b04c3aa0aad2870cfe799f32f8b59789c98e1816bbce9e83f4823c5b840b/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-win32.whl", hash = "sha256:fca724a21a372731edb290841edd28a9fb1ee490f833392752844ac807c0086a", size = 552682, upload-time = "2026-01-28T05:58:05.649Z" }, { url = "https://files.pythonhosted.org/packages/bd/e1/6d6816b296a529ac9b897ad228b1e084eb1f92319e96371880eebdc874a6/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-win_amd64.whl", hash = "sha256:823c0bd7770977d4b10e0ed0aef2f3682276b7c88b8b65cfc540afce5951392f", size = 559464, upload-time = "2026-01-28T05:58:07.261Z" }, + { url = "https://files.pythonhosted.org/packages/8c/a8/d4dab8a58fc2e6981fc7a58c4e56ba9d777fb24931cec6a22152edbb3540/glfw-2.10.0-py2.py3-none-macosx_10_6_intel.whl", hash = "sha256:a0d1f29f206219cc291edfb6cace663a86da2470632551c998e3db82d48ea177", size = 105288, upload-time = "2026-03-10T17:21:19.929Z" }, + { url = "https://files.pythonhosted.org/packages/14/61/68d35e001872a7705112418da236fa2418d4f2e5419f8b2837f9b81bb3da/glfw-2.10.0-py2.py3-none-macosx_11_0_arm64.whl", hash = "sha256:d28d6f3ef217e64e35dc6fd0a7acb4cec9bfe7cd14dd9b35a7228a87002de154", size = 102139, upload-time = "2026-03-10T17:21:21.645Z" }, + { url = "https://files.pythonhosted.org/packages/4e/e1/ca5984081aaae07c9d371cb11dc4e4ff603510678ed9b73e58b6c351fe63/glfw-2.10.0-py2.py3-none-manylinux2014_aarch64.whl", hash = "sha256:f968b522bb6a0e04aaf4dcac30a476d7229308bb2bac406a60587debb5a61e29", size = 229998, upload-time = "2026-03-10T17:21:23.549Z" }, + { url = "https://files.pythonhosted.org/packages/fa/c4/82ac75fdcfba2896da7a573c0fc7f8ceb8f77ead6866d500d06c32f1c464/glfw-2.10.0-py2.py3-none-manylinux2014_x86_64.whl", hash = "sha256:68cf3752bdadb6f4bc0a876247c28c88c7251ac39f8af076ed938fdfd71e72dd", size = 241944, upload-time = "2026-03-10T17:21:26.102Z" }, + { url = "https://files.pythonhosted.org/packages/e3/96/9f691823cca5eb6a08f346bd0ff03b78032db9370b509a1e9c8976fb20a5/glfw-2.10.0-py2.py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:44d98de5dbf8f727e0cb29f9b29d29528ea7570f2e6f42f8430a69df05f12b48", size = 231009, upload-time = "2026-03-10T17:21:28.481Z" }, + { url = "https://files.pythonhosted.org/packages/3f/93/977b9e679e356871d428ae7a1139ec767dd5177bed58a6344b4d2199e00f/glfw-2.10.0-py2.py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:cca5158d62189e08792b1ae54f92307a282921a0e7783315b467e21b0a381c88", size = 243480, upload-time = "2026-03-10T17:21:30.538Z" }, + { url = "https://files.pythonhosted.org/packages/f9/bd/cea9569c8f2188b0a104472951420434a3e1f5cf26f5836ef9d7227a1a30/glfw-2.10.0-py2.py3-none-win32.whl", hash = "sha256:5e024509989740e8e7b86cc4aab508195495f79879072b0e1f68bd036a2916ad", size = 552641, upload-time = "2026-03-10T17:21:32.653Z" }, + { url = "https://files.pythonhosted.org/packages/cc/9b/4366ad3e1c0688146c70aa6143584d6a8d88583b9390f106250e25a3d5cd/glfw-2.10.0-py2.py3-none-win_amd64.whl", hash = "sha256:7f787ee8645781f10e8800438ce4357ab38c573ffb191aba380c1e72eba6311c", size = 559423, upload-time = "2026-03-10T17:21:34.766Z" }, +] + +[[package]] +name = "grpcio" +version = "1.78.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/8a/3d098f35c143a89520e568e6539cc098fcd294495910e359889ce8741c84/grpcio-1.78.0.tar.gz", hash = "sha256:7382b95189546f375c174f53a5fa873cef91c4b8005faa05cc5b3beea9c4f1c5", size = 12852416, upload-time = "2026-02-06T09:57:18.093Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4e/f4/7384ed0178203d6074446b3c4f46c90a22ddf7ae0b3aee521627f54cfc2a/grpcio-1.78.0-cp312-cp312-linux_armv7l.whl", hash = "sha256:f9ab915a267fc47c7e88c387a3a28325b58c898e23d4995f765728f4e3dedb97", size = 5913985, upload-time = "2026-02-06T09:55:26.832Z" }, + { url = "https://files.pythonhosted.org/packages/81/ed/be1caa25f06594463f685b3790b320f18aea49b33166f4141bfdc2bfb236/grpcio-1.78.0-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:3f8904a8165ab21e07e58bf3e30a73f4dffc7a1e0dbc32d51c61b5360d26f43e", size = 11811853, upload-time = "2026-02-06T09:55:29.224Z" }, + { url = "https://files.pythonhosted.org/packages/24/a7/f06d151afc4e64b7e3cc3e872d331d011c279aaab02831e40a81c691fb65/grpcio-1.78.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:859b13906ce098c0b493af92142ad051bf64c7870fa58a123911c88606714996", size = 6475766, upload-time = "2026-02-06T09:55:31.825Z" }, + { url = "https://files.pythonhosted.org/packages/8a/a8/4482922da832ec0082d0f2cc3a10976d84a7424707f25780b82814aafc0a/grpcio-1.78.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:b2342d87af32790f934a79c3112641e7b27d63c261b8b4395350dad43eff1dc7", size = 7170027, upload-time = "2026-02-06T09:55:34.7Z" }, + { url = "https://files.pythonhosted.org/packages/54/bf/f4a3b9693e35d25b24b0b39fa46d7d8a3c439e0a3036c3451764678fec20/grpcio-1.78.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:12a771591ae40bc65ba67048fa52ef4f0e6db8279e595fd349f9dfddeef571f9", size = 6690766, upload-time = "2026-02-06T09:55:36.902Z" }, + { url = "https://files.pythonhosted.org/packages/c7/b9/521875265cc99fe5ad4c5a17010018085cae2810a928bf15ebe7d8bcd9cc/grpcio-1.78.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:185dea0d5260cbb2d224c507bf2a5444d5abbb1fa3594c1ed7e4c709d5eb8383", size = 7266161, upload-time = "2026-02-06T09:55:39.824Z" }, + { url = "https://files.pythonhosted.org/packages/05/86/296a82844fd40a4ad4a95f100b55044b4f817dece732bf686aea1a284147/grpcio-1.78.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:51b13f9aed9d59ee389ad666b8c2214cc87b5de258fa712f9ab05f922e3896c6", size = 8253303, upload-time = "2026-02-06T09:55:42.353Z" }, + { url = "https://files.pythonhosted.org/packages/f3/e4/ea3c0caf5468537f27ad5aab92b681ed7cc0ef5f8c9196d3fd42c8c2286b/grpcio-1.78.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fd5f135b1bd58ab088930b3c613455796dfa0393626a6972663ccdda5b4ac6ce", size = 7698222, upload-time = "2026-02-06T09:55:44.629Z" }, + { url = "https://files.pythonhosted.org/packages/d7/47/7f05f81e4bb6b831e93271fb12fd52ba7b319b5402cbc101d588f435df00/grpcio-1.78.0-cp312-cp312-win32.whl", hash = "sha256:94309f498bcc07e5a7d16089ab984d42ad96af1d94b5a4eb966a266d9fcabf68", size = 4066123, upload-time = "2026-02-06T09:55:47.644Z" }, + { url = "https://files.pythonhosted.org/packages/ad/e7/d6914822c88aa2974dbbd10903d801a28a19ce9cd8bad7e694cbbcf61528/grpcio-1.78.0-cp312-cp312-win_amd64.whl", hash = "sha256:9566fe4ababbb2610c39190791e5b829869351d14369603702e890ef3ad2d06e", size = 4797657, upload-time = "2026-02-06T09:55:49.86Z" }, +] + +[[package]] +name = "gym" +version = "0.26.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cloudpickle" }, + { name = "gym-notices" }, + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ab/b1/eb05a423eb801ab7d0715d6a3b28d92589e30b437052553df19ca2087240/gym-0.26.2.tar.gz", hash = "sha256:e0d882f4b54f0c65f203104c24ab8a38b039f1289986803c7d02cdbe214fbcc4", size = 721689, upload-time = "2022-10-04T23:57:43.247Z" } + +[[package]] +name = "gym-notices" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a2/4d/035922b950b224ee4b65a9a4550a22eac8985a3f0e1ef42546d9047e7a72/gym_notices-0.1.0.tar.gz", hash = "sha256:9f9477ef68a8c15e42625d4fa53631237e3e6ae947f325b5c149c081499adc1b", size = 3084, upload-time = "2025-07-27T10:12:41.534Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/41/55/55d157aa8693090954fc9639bf27218240517c3bc7afa6e97412da6ebfd9/gym_notices-0.1.0-py3-none-any.whl", hash = "sha256:a943af4446cb619d04fd1e470b9272b4473e08a06d1c7cc9005755a4a0b8c905", size = 3349, upload-time = "2025-07-27T10:12:40.039Z" }, ] [[package]] @@ -537,6 +756,29 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c5/7b/bca5613a0c3b542420cf92bd5e5fb8ebd5435ce1011a091f66bb7693285e/humanize-4.15.0-py3-none-any.whl", hash = "sha256:b1186eb9f5a9749cd9cb8565aee77919dd7c8d076161cf44d70e59e3301e1769", size = 132203, upload-time = "2025-12-20T20:16:11.67Z" }, ] +[[package]] +name = "hydra-core" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "omegaconf" }, + { name = "packaging" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6d/8e/07e42bc434a847154083b315779b0a81d567154504624e181caf2c71cd98/hydra-core-1.3.2.tar.gz", hash = "sha256:8a878ed67216997c3e9d88a8e72e7b4767e81af37afb4ea3334b269a4390a824", size = 3263494, upload-time = "2023-02-23T18:33:43.03Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c6/50/e0edd38dcd63fb26a8547f13d28f7a008bc4a3fd4eb4ff030673f22ad41a/hydra_core-1.3.2-py3-none-any.whl", hash = "sha256:fa0238a9e31df3373b35b0bfb672c34cc92718d21f81311d8996a16de1141d8b", size = 154547, upload-time = "2023-02-23T18:33:40.801Z" }, +] + +[[package]] +name = "identify" +version = "2.6.18" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/46/c4/7fb4db12296cdb11893d61c92048fe617ee853f8523b9b296ac03b43757e/identify-2.6.18.tar.gz", hash = "sha256:873ac56a5e3fd63e7438a7ecbc4d91aca692eb3fefa4534db2b7913f3fc352fd", size = 99580, upload-time = "2026-03-15T18:39:50.319Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/33/92ef41c6fad0233e41d3d84ba8e8ad18d1780f1e5d99b3c683e6d7f98b63/identify-2.6.18-py2.py3-none-any.whl", hash = "sha256:8db9d3c8ea9079db92cafb0ebf97abdc09d52e97f4dcf773a2e694048b7cd737", size = 99394, upload-time = "2026-03-15T18:39:48.915Z" }, +] + [[package]] name = "idna" version = "3.11" @@ -548,24 +790,38 @@ wheels = [ [[package]] name = "imageio" -version = "2.37.2" +version = "2.37.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, { name = "pillow" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a3/6f/606be632e37bf8d05b253e8626c2291d74c691ddc7bcdf7d6aaf33b32f6a/imageio-2.37.2.tar.gz", hash = "sha256:0212ef2727ac9caa5ca4b2c75ae89454312f440a756fcfc8ef1993e718f50f8a", size = 389600, upload-time = "2025-11-04T14:29:39.898Z" } +sdist = { url = "https://files.pythonhosted.org/packages/b1/84/93bcd1300216ea50811cee96873b84a1bebf8d0489ffaf7f2a3756bab866/imageio-2.37.3.tar.gz", hash = "sha256:bbb37efbfc4c400fcd534b367b91fcd66d5da639aaa138034431a1c5e0a41451", size = 389673, upload-time = "2026-03-09T11:31:12.573Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/fb/fe/301e0936b79bcab4cacc7548bf2853fc28dced0a578bab1f7ef53c9aa75b/imageio-2.37.2-py3-none-any.whl", hash = "sha256:ad9adfb20335d718c03de457358ed69f141021a333c40a53e57273d8a5bd0b9b", size = 317646, upload-time = "2025-11-04T14:29:37.948Z" }, + { url = "https://files.pythonhosted.org/packages/49/fa/391e437a34e55095173dca5f24070d89cbc233ff85bf1c29c93248c6588d/imageio-2.37.3-py3-none-any.whl", hash = "sha256:46f5bb8522cd421c0f5ae104d8268f569d856b29eb1a13b92829d1970f32c9f0", size = 317646, upload-time = "2026-03-09T11:31:10.771Z" }, ] [[package]] -name = "importlib-resources" -version = "6.5.2" +name = "imageio-ffmpeg" +version = "0.6.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/cf/8c/f834fbf984f691b4f7ff60f50b514cc3de5cc08abfc3295564dd89c5e2e7/importlib_resources-6.5.2.tar.gz", hash = "sha256:185f87adef5bcc288449d98fb4fba07cea78bc036455dd44c5fc4a2fe78fed2c", size = 44693, upload-time = "2025-01-03T18:51:56.698Z" } +sdist = { url = "https://files.pythonhosted.org/packages/44/bd/c3343c721f2a1b0c9fc71c1aebf1966a3b7f08c2eea8ed5437a2865611d6/imageio_ffmpeg-0.6.0.tar.gz", hash = "sha256:e2556bed8e005564a9f925bb7afa4002d82770d6b08825078b7697ab88ba1755", size = 25210, upload-time = "2025-01-16T21:34:32.747Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a4/ed/1f1afb2e9e7f38a545d628f864d562a5ae64fe6f7a10e28ffb9b185b4e89/importlib_resources-6.5.2-py3-none-any.whl", hash = "sha256:789cfdc3ed28c78b67a06acb8126751ced69a3d5f79c095a98298cd8a760ccec", size = 37461, upload-time = "2025-01-03T18:51:54.306Z" }, + { url = "https://files.pythonhosted.org/packages/da/58/87ef68ac83f4c7690961bce288fd8e382bc5f1513860fc7f90a9c1c1c6bf/imageio_ffmpeg-0.6.0-py3-none-macosx_10_9_intel.macosx_10_9_x86_64.whl", hash = "sha256:9d2baaf867088508d4a3458e61eeb30e945c4ad8016025545f66c4b5aaef0a61", size = 24932969, upload-time = "2025-01-16T21:34:20.464Z" }, + { url = "https://files.pythonhosted.org/packages/40/5c/f3d8a657d362cc93b81aab8feda487317da5b5d31c0e1fdfd5e986e55d17/imageio_ffmpeg-0.6.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:b1ae3173414b5fc5f538a726c4e48ea97edc0d2cdc11f103afee655c463fa742", size = 21113891, upload-time = "2025-01-16T21:34:00.277Z" }, + { url = "https://files.pythonhosted.org/packages/33/e7/1925bfbc563c39c1d2e82501d8372734a5c725e53ac3b31b4c2d081e895b/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:1d47bebd83d2c5fc770720d211855f208af8a596c82d17730aa51e815cdee6dc", size = 25632706, upload-time = "2025-01-16T21:33:53.475Z" }, + { url = "https://files.pythonhosted.org/packages/a0/2d/43c8522a2038e9d0e7dbdf3a61195ecc31ca576fb1527a528c877e87d973/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:c7e46fcec401dd990405049d2e2f475e2b397779df2519b544b8aab515195282", size = 29498237, upload-time = "2025-01-16T21:34:13.726Z" }, + { url = "https://files.pythonhosted.org/packages/a0/13/59da54728351883c3c1d9fca1710ab8eee82c7beba585df8f25ca925f08f/imageio_ffmpeg-0.6.0-py3-none-win32.whl", hash = "sha256:196faa79366b4a82f95c0f4053191d2013f4714a715780f0ad2a68ff37483cc2", size = 19652251, upload-time = "2025-01-16T21:34:06.812Z" }, + { url = "https://files.pythonhosted.org/packages/2c/c6/fa760e12a2483469e2bf5058c5faff664acf66cadb4df2ad6205b016a73d/imageio_ffmpeg-0.6.0-py3-none-win_amd64.whl", hash = "sha256:02fa47c83703c37df6bfe4896aab339013f62bf02c5ebf2dce6da56af04ffc0a", size = 31246824, upload-time = "2025-01-16T21:34:28.6Z" }, +] + +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, ] [[package]] @@ -594,7 +850,7 @@ wheels = [ [[package]] name = "ipython" -version = "9.10.0" +version = "9.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "colorama", marker = "sys_platform == 'win32'" }, @@ -608,9 +864,9 @@ dependencies = [ { name = "stack-data" }, { name = "traitlets" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a6/60/2111715ea11f39b1535bed6024b7dec7918b71e5e5d30855a5b503056b50/ipython-9.10.0.tar.gz", hash = "sha256:cd9e656be97618a0676d058134cd44e6dc7012c0e5cb36a9ce96a8c904adaf77", size = 4426526, upload-time = "2026-02-02T10:00:33.594Z" } +sdist = { url = "https://files.pythonhosted.org/packages/86/28/a4698eda5a8928a45d6b693578b135b753e14fa1c2b36ee9441e69a45576/ipython-9.11.0.tar.gz", hash = "sha256:2a94bc4406b22ecc7e4cb95b98450f3ea493a76bec8896cda11b78d7752a6667", size = 4427354, upload-time = "2026-03-05T08:57:30.549Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3d/aa/898dec789a05731cd5a9f50605b7b44a72bd198fd0d4528e11fc610177cc/ipython-9.10.0-py3-none-any.whl", hash = "sha256:c6ab68cc23bba8c7e18e9b932797014cc61ea7fd6f19de180ab9ba73e65ee58d", size = 622774, upload-time = "2026-02-02T10:00:31.503Z" }, + { url = "https://files.pythonhosted.org/packages/b2/90/45c72becc57158facc6a6404f663b77bbcea2519ca57f760e2879ae1315d/ipython-9.11.0-py3-none-any.whl", hash = "sha256:6922d5bcf944c6e525a76a0a304451b60a2b6f875e86656d8bc2dfda5d710e19", size = 624222, upload-time = "2026-03-05T08:57:28.94Z" }, ] [[package]] @@ -701,6 +957,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bb/02/265e5ccadd65fee2f0716431573d9e512e5c6aecb23f478a7a92053cf219/jaxlib-0.9.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:08733d1431238a7cf9108338ab7be898b97181cba0eef53f2f9fd3de17d20adb", size = 60508788, upload-time = "2026-02-05T18:46:43.209Z" }, ] +[[package]] +name = "jaxtyping" +version = "0.3.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wadler-lindig" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c2/be/00294e369938937e31b094437d5ea040e4fd1a20b998ebe572c4a1dcfa68/jaxtyping-0.3.9.tar.gz", hash = "sha256:f8c02d1b623d5f1b6665d4f3ddaec675d70004f16a792102c2fc51264190951d", size = 45857, upload-time = "2026-02-16T10:35:13.263Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/94/05/3e39d416fb92b2738a76e8265e6bfc5d10542f90a7c32ad1eb831eea3fa3/jaxtyping-0.3.9-py3-none-any.whl", hash = "sha256:a00557a9d616eff157491f06ed2e21ed94886fad3832399273eb912b345da378", size = 56274, upload-time = "2026-02-16T10:35:11.795Z" }, +] + [[package]] name = "jedi" version = "0.19.2" @@ -713,6 +981,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c0/5a/9cac0c82afec3d09ccd97c8b6502d48f165f9124db81b4bcb90b4af974ee/jedi-0.19.2-py2.py3-none-any.whl", hash = "sha256:a8ef22bde8490f57fe5c7681a3c83cb58874daf72b4784de3cce5b6ef6edb5b9", size = 1572278, upload-time = "2024-11-11T01:41:40.175Z" }, ] +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, +] + [[package]] name = "jupyter-client" version = "8.8.0" @@ -744,23 +1024,29 @@ wheels = [ [[package]] name = "kiwisolver" -version = "1.4.9" +version = "1.5.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/5c/3c/85844f1b0feb11ee581ac23fe5fce65cd049a200c1446708cc1b7f922875/kiwisolver-1.4.9.tar.gz", hash = "sha256:c3b22c26c6fd6811b0ae8363b95ca8ce4ea3c202d3d0975b2914310ceb1bcc4d", size = 97564, upload-time = "2025-08-10T21:27:49.279Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d0/67/9c61eccb13f0bdca9307614e782fec49ffdde0f7a2314935d489fa93cd9c/kiwisolver-1.5.0.tar.gz", hash = "sha256:d4193f3d9dc3f6f79aaed0e5637f45d98850ebf01f7ca20e69457f3e8946b66a", size = 103482, upload-time = "2026-03-09T13:15:53.382Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/86/c9/13573a747838aeb1c76e3267620daa054f4152444d1f3d1a2324b78255b5/kiwisolver-1.4.9-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:ac5a486ac389dddcc5bef4f365b6ae3ffff2c433324fb38dd35e3fab7c957999", size = 123686, upload-time = "2025-08-10T21:26:10.034Z" }, - { url = "https://files.pythonhosted.org/packages/51/ea/2ecf727927f103ffd1739271ca19c424d0e65ea473fbaeea1c014aea93f6/kiwisolver-1.4.9-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:f2ba92255faa7309d06fe44c3a4a97efe1c8d640c2a79a5ef728b685762a6fd2", size = 66460, upload-time = "2025-08-10T21:26:11.083Z" }, - { url = "https://files.pythonhosted.org/packages/5b/5a/51f5464373ce2aeb5194508298a508b6f21d3867f499556263c64c621914/kiwisolver-1.4.9-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4a2899935e724dd1074cb568ce7ac0dce28b2cd6ab539c8e001a8578eb106d14", size = 64952, upload-time = "2025-08-10T21:26:12.058Z" }, - { url = "https://files.pythonhosted.org/packages/70/90/6d240beb0f24b74371762873e9b7f499f1e02166a2d9c5801f4dbf8fa12e/kiwisolver-1.4.9-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f6008a4919fdbc0b0097089f67a1eb55d950ed7e90ce2cc3e640abadd2757a04", size = 1474756, upload-time = "2025-08-10T21:26:13.096Z" }, - { url = "https://files.pythonhosted.org/packages/12/42/f36816eaf465220f683fb711efdd1bbf7a7005a2473d0e4ed421389bd26c/kiwisolver-1.4.9-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:67bb8b474b4181770f926f7b7d2f8c0248cbcb78b660fdd41a47054b28d2a752", size = 1276404, upload-time = "2025-08-10T21:26:14.457Z" }, - { url = "https://files.pythonhosted.org/packages/2e/64/bc2de94800adc830c476dce44e9b40fd0809cddeef1fde9fcf0f73da301f/kiwisolver-1.4.9-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:2327a4a30d3ee07d2fbe2e7933e8a37c591663b96ce42a00bc67461a87d7df77", size = 1294410, upload-time = "2025-08-10T21:26:15.73Z" }, - { url = "https://files.pythonhosted.org/packages/5f/42/2dc82330a70aa8e55b6d395b11018045e58d0bb00834502bf11509f79091/kiwisolver-1.4.9-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:7a08b491ec91b1d5053ac177afe5290adacf1f0f6307d771ccac5de30592d198", size = 1343631, upload-time = "2025-08-10T21:26:17.045Z" }, - { url = "https://files.pythonhosted.org/packages/22/fd/f4c67a6ed1aab149ec5a8a401c323cee7a1cbe364381bb6c9c0d564e0e20/kiwisolver-1.4.9-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:d8fc5c867c22b828001b6a38d2eaeb88160bf5783c6cb4a5e440efc981ce286d", size = 2224963, upload-time = "2025-08-10T21:26:18.737Z" }, - { url = "https://files.pythonhosted.org/packages/45/aa/76720bd4cb3713314677d9ec94dcc21ced3f1baf4830adde5bb9b2430a5f/kiwisolver-1.4.9-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:3b3115b2581ea35bb6d1f24a4c90af37e5d9b49dcff267eeed14c3893c5b86ab", size = 2321295, upload-time = "2025-08-10T21:26:20.11Z" }, - { url = "https://files.pythonhosted.org/packages/80/19/d3ec0d9ab711242f56ae0dc2fc5d70e298bb4a1f9dfab44c027668c673a1/kiwisolver-1.4.9-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:858e4c22fb075920b96a291928cb7dea5644e94c0ee4fcd5af7e865655e4ccf2", size = 2487987, upload-time = "2025-08-10T21:26:21.49Z" }, - { url = "https://files.pythonhosted.org/packages/39/e9/61e4813b2c97e86b6fdbd4dd824bf72d28bcd8d4849b8084a357bc0dd64d/kiwisolver-1.4.9-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:ed0fecd28cc62c54b262e3736f8bb2512d8dcfdc2bcf08be5f47f96bf405b145", size = 2291817, upload-time = "2025-08-10T21:26:22.812Z" }, - { url = "https://files.pythonhosted.org/packages/a0/41/85d82b0291db7504da3c2defe35c9a8a5c9803a730f297bd823d11d5fb77/kiwisolver-1.4.9-cp312-cp312-win_amd64.whl", hash = "sha256:f68208a520c3d86ea51acf688a3e3002615a7f0238002cccc17affecc86a8a54", size = 73895, upload-time = "2025-08-10T21:26:24.37Z" }, - { url = "https://files.pythonhosted.org/packages/e2/92/5f3068cf15ee5cb624a0c7596e67e2a0bb2adee33f71c379054a491d07da/kiwisolver-1.4.9-cp312-cp312-win_arm64.whl", hash = "sha256:2c1a4f57df73965f3f14df20b80ee29e6a7930a57d2d9e8491a25f676e197c60", size = 64992, upload-time = "2025-08-10T21:26:25.732Z" }, + { url = "https://files.pythonhosted.org/packages/4d/b2/818b74ebea34dabe6d0c51cb1c572e046730e64844da6ed646d5298c40ce/kiwisolver-1.5.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:4e9750bc21b886308024f8a54ccb9a2cc38ac9fa813bf4348434e3d54f337ff9", size = 123158, upload-time = "2026-03-09T13:13:23.127Z" }, + { url = "https://files.pythonhosted.org/packages/bf/d9/405320f8077e8e1c5c4bd6adc45e1e6edf6d727b6da7f2e2533cf58bff71/kiwisolver-1.5.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:72ec46b7eba5b395e0a7b63025490d3214c11013f4aacb4f5e8d6c3041829588", size = 66388, upload-time = "2026-03-09T13:13:24.765Z" }, + { url = "https://files.pythonhosted.org/packages/99/9f/795fedf35634f746151ca8839d05681ceb6287fbed6cc1c9bf235f7887c2/kiwisolver-1.5.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ed3a984b31da7481b103f68776f7128a89ef26ed40f4dc41a2223cda7fb24819", size = 64068, upload-time = "2026-03-09T13:13:25.878Z" }, + { url = "https://files.pythonhosted.org/packages/c4/13/680c54afe3e65767bed7ec1a15571e1a2f1257128733851ade24abcefbcc/kiwisolver-1.5.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bb5136fb5352d3f422df33f0c879a1b0c204004324150cc3b5e3c4f310c9049f", size = 1477934, upload-time = "2026-03-09T13:13:27.166Z" }, + { url = "https://files.pythonhosted.org/packages/c8/2f/cebfcdb60fd6a9b0f6b47a9337198bcbad6fbe15e68189b7011fd914911f/kiwisolver-1.5.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b2af221f268f5af85e776a73d62b0845fc8baf8ef0abfae79d29c77d0e776aaf", size = 1278537, upload-time = "2026-03-09T13:13:28.707Z" }, + { url = "https://files.pythonhosted.org/packages/f2/0d/9b782923aada3fafb1d6b84e13121954515c669b18af0c26e7d21f579855/kiwisolver-1.5.0-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:b0f172dc8ffaccb8522d7c5d899de00133f2f1ca7b0a49b7da98e901de87bf2d", size = 1296685, upload-time = "2026-03-09T13:13:30.528Z" }, + { url = "https://files.pythonhosted.org/packages/27/70/83241b6634b04fe44e892688d5208332bde130f38e610c0418f9ede47ded/kiwisolver-1.5.0-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:6ab8ba9152203feec73758dad83af9a0bbe05001eb4639e547207c40cfb52083", size = 1346024, upload-time = "2026-03-09T13:13:32.818Z" }, + { url = "https://files.pythonhosted.org/packages/e4/db/30ed226fb271ae1a6431fc0fe0edffb2efe23cadb01e798caeb9f2ceae8f/kiwisolver-1.5.0-cp312-cp312-manylinux_2_39_riscv64.whl", hash = "sha256:cdee07c4d7f6d72008d3f73b9bf027f4e11550224c7c50d8df1ae4a37c1402a6", size = 987241, upload-time = "2026-03-09T13:13:34.435Z" }, + { url = "https://files.pythonhosted.org/packages/ec/bd/c314595208e4c9587652d50959ead9e461995389664e490f4dce7ff0f782/kiwisolver-1.5.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:7c60d3c9b06fb23bd9c6139281ccbdc384297579ae037f08ae90c69f6845c0b1", size = 2227742, upload-time = "2026-03-09T13:13:36.4Z" }, + { url = "https://files.pythonhosted.org/packages/c1/43/0499cec932d935229b5543d073c2b87c9c22846aab48881e9d8d6e742a2d/kiwisolver-1.5.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:e315e5ec90d88e140f57696ff85b484ff68bb311e36f2c414aa4286293e6dee0", size = 2323966, upload-time = "2026-03-09T13:13:38.204Z" }, + { url = "https://files.pythonhosted.org/packages/3d/6f/79b0d760907965acfd9d61826a3d41f8f093c538f55cd2633d3f0db269f6/kiwisolver-1.5.0-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:1465387ac63576c3e125e5337a6892b9e99e0627d52317f3ca79e6930d889d15", size = 1977417, upload-time = "2026-03-09T13:13:39.966Z" }, + { url = "https://files.pythonhosted.org/packages/ab/31/01d0537c41cb75a551a438c3c7a80d0c60d60b81f694dac83dd436aec0d0/kiwisolver-1.5.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:530a3fd64c87cffa844d4b6b9768774763d9caa299e9b75d8eca6a4423b31314", size = 2491238, upload-time = "2026-03-09T13:13:41.698Z" }, + { url = "https://files.pythonhosted.org/packages/e4/34/8aefdd0be9cfd00a44509251ba864f5caf2991e36772e61c408007e7f417/kiwisolver-1.5.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:1d9daea4ea6b9be74fe2f01f7fbade8d6ffab263e781274cffca0dba9be9eec9", size = 2294947, upload-time = "2026-03-09T13:13:43.343Z" }, + { url = "https://files.pythonhosted.org/packages/ad/cf/0348374369ca588f8fe9c338fae49fa4e16eeb10ffb3d012f23a54578a9e/kiwisolver-1.5.0-cp312-cp312-win_amd64.whl", hash = "sha256:f18c2d9782259a6dc132fdc7a63c168cbc74b35284b6d75c673958982a378384", size = 73569, upload-time = "2026-03-09T13:13:45.792Z" }, + { url = "https://files.pythonhosted.org/packages/28/26/192b26196e2316e2bd29deef67e37cdf9870d9af8e085e521afff0fed526/kiwisolver-1.5.0-cp312-cp312-win_arm64.whl", hash = "sha256:f7c7553b13f69c1b29a5bde08ddc6d9d0c8bfb84f9ed01c30db25944aeb852a7", size = 64997, upload-time = "2026-03-09T13:13:46.878Z" }, + { url = "https://files.pythonhosted.org/packages/1c/fa/2910df836372d8761bb6eff7d8bdcb1613b5c2e03f260efe7abe34d388a7/kiwisolver-1.5.0-graalpy312-graalpy250_312_native-macosx_10_13_x86_64.whl", hash = "sha256:5ae8e62c147495b01a0f4765c878e9bfdf843412446a247e28df59936e99e797", size = 130262, upload-time = "2026-03-09T13:15:35.629Z" }, + { url = "https://files.pythonhosted.org/packages/0f/41/c5f71f9f00aabcc71fee8b7475e3f64747282580c2fe748961ba29b18385/kiwisolver-1.5.0-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:f6764a4ccab3078db14a632420930f6186058750df066b8ea2a7106df91d3203", size = 138036, upload-time = "2026-03-09T13:15:36.894Z" }, + { url = "https://files.pythonhosted.org/packages/fa/06/7399a607f434119c6e1fdc8ec89a8d51ccccadf3341dee4ead6bd14caaf5/kiwisolver-1.5.0-graalpy312-graalpy250_312_native-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c31c13da98624f957b0fb1b5bae5383b2333c2c3f6793d9825dd5ce79b525cb7", size = 194295, upload-time = "2026-03-09T13:15:38.22Z" }, + { url = "https://files.pythonhosted.org/packages/b5/91/53255615acd2a1eaca307ede3c90eb550bae9c94581f8c00081b6b1c8f44/kiwisolver-1.5.0-graalpy312-graalpy250_312_native-win_amd64.whl", hash = "sha256:1f1489f769582498610e015a8ef2d36f28f505ab3096d0e16b4858a9ec214f57", size = 75987, upload-time = "2026-03-09T13:15:39.65Z" }, ] [[package]] @@ -807,6 +1093,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ea/7b/93c73c67db235931527301ed3785f849c78991e2e34f3fd9a6663ffda4c5/lxml-6.0.2-cp312-cp312-win_arm64.whl", hash = "sha256:61cb10eeb95570153e0c0e554f58df92ecf5109f75eacad4a95baa709e26c3d6", size = 3672836, upload-time = "2025-09-22T04:01:52.145Z" }, ] +[[package]] +name = "markdown" +version = "3.10.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2b/f4/69fa6ed85ae003c2378ffa8f6d2e3234662abd02c10d216c0ba96081a238/markdown-3.10.2.tar.gz", hash = "sha256:994d51325d25ad8aa7ce4ebaec003febcce822c3f8c911e3b17c52f7f589f950", size = 368805, upload-time = "2026-02-09T14:57:26.942Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/de/1f/77fa3081e4f66ca3576c896ae5d31c3002ac6607f9747d2e3aa49227e464/markdown-3.10.2-py3-none-any.whl", hash = "sha256:e91464b71ae3ee7afd3017d9f358ef0baf158fd9a298db92f1d4761133824c36", size = 108180, upload-time = "2026-02-09T14:57:25.787Z" }, +] + [[package]] name = "markdown-it-py" version = "4.0.0" @@ -819,6 +1114,37 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl", hash = "sha256:87327c59b172c5011896038353a81343b6754500a08cd7a4973bb48c6d578147", size = 87321, upload-time = "2025-08-11T12:57:51.923Z" }, ] +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e", size = 11615, upload-time = "2025-09-27T18:36:30.854Z" }, + { url = "https://files.pythonhosted.org/packages/9a/81/7e4e08678a1f98521201c3079f77db69fb552acd56067661f8c2f534a718/markupsafe-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce", size = 12020, upload-time = "2025-09-27T18:36:31.971Z" }, + { url = "https://files.pythonhosted.org/packages/1e/2c/799f4742efc39633a1b54a92eec4082e4f815314869865d876824c257c1e/markupsafe-3.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d", size = 24332, upload-time = "2025-09-27T18:36:32.813Z" }, + { url = "https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" }, + { url = "https://files.pythonhosted.org/packages/2c/54/887f3092a85238093a0b2154bd629c89444f395618842e8b0c41783898ea/markupsafe-3.0.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a", size = 21962, upload-time = "2025-09-27T18:36:35.099Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2f/336b8c7b6f4a4d95e91119dc8521402461b74a485558d8f238a68312f11c/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b", size = 23760, upload-time = "2025-09-27T18:36:36.001Z" }, + { url = "https://files.pythonhosted.org/packages/32/43/67935f2b7e4982ffb50a4d169b724d74b62a3964bc1a9a527f5ac4f1ee2b/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f", size = 21529, upload-time = "2025-09-27T18:36:36.906Z" }, + { url = "https://files.pythonhosted.org/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" }, + { url = "https://files.pythonhosted.org/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d", size = 14540, upload-time = "2025-09-27T18:36:38.761Z" }, + { url = "https://files.pythonhosted.org/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c", size = 15105, upload-time = "2025-09-27T18:36:39.701Z" }, + { url = "https://files.pythonhosted.org/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, +] + +[[package]] +name = "marshmallow" +version = "3.26.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "packaging" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/55/79/de6c16cc902f4fc372236926b0ce2ab7845268dcc30fb2fbb7f71b418631/marshmallow-3.26.2.tar.gz", hash = "sha256:bbe2adb5a03e6e3571b573f42527c6fe926e17467833660bebd11593ab8dfd57", size = 222095, upload-time = "2025-12-22T06:53:53.309Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/be/2f/5108cb3ee4ba6501748c4908b908e55f42a5b66245b4cfe0c99326e1ef6e/marshmallow-3.26.2-py3-none-any.whl", hash = "sha256:013fa8a3c4c276c24d26d84ce934dc964e2aa794345a0f8c7e5a7191482c8a73", size = 50964, upload-time = "2025-12-22T06:53:51.801Z" }, +] + [[package]] name = "matplotlib" version = "3.10.8" @@ -919,6 +1245,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/72/41f8cee40c465fa6f73ef984a88ece724b3e8731eb3efdaf483863b54422/moojoco-1.1.7-py3-none-any.whl", hash = "sha256:8f8bb333fc12e94d527ce9f4365e485af8bdb93a3f5257be568fcf8e1910a43f", size = 24405, upload-time = "2025-05-09T10:37:04.242Z" }, ] +[[package]] +name = "mpmath" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, +] + [[package]] name = "msgpack" version = "1.1.2" @@ -938,7 +1273,7 @@ wheels = [ [[package]] name = "mujoco" -version = "3.5.0" +version = "3.6.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "absl-py" }, @@ -947,18 +1282,18 @@ dependencies = [ { name = "numpy" }, { name = "pyopengl" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/56/0d/005f0d49ad5878f0611a7c018550b8504d480a7a17ad7e6773ff47d8627a/mujoco-3.5.0.tar.gz", hash = "sha256:5c85a6fc7560ab5fa4534f35ff459e12dc3609681f307e457dbb49b6217f4d73", size = 912543, upload-time = "2026-02-13T01:02:51.554Z" } +sdist = { url = "https://files.pythonhosted.org/packages/42/82/f8f08dfe9123df4351b560f894f0e7166c1a45a0dd2f04145ed00b8f849b/mujoco-3.6.0.tar.gz", hash = "sha256:15c89f423e33bce0860ad7061763b72323426d6348d7b2e46ebdcc37b11e0905", size = 915041, upload-time = "2026-03-11T01:45:42.807Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f9/f0/4772421643f1c5aaf46d9e500a8716f59b02c8bf30bfa92cb8a763159efb/mujoco-3.5.0-cp312-cp312-macosx_10_16_x86_64.whl", hash = "sha256:ec0587cc423385a8d45343a981df58511cb69758ba99164a71567af2d41be3c9", size = 7100581, upload-time = "2026-02-13T01:02:29.182Z" }, - { url = "https://files.pythonhosted.org/packages/e1/d4/d0032323f58a9b8080b8464c6aade8d5ac2e101dbed1de64a38b3913b446/mujoco-3.5.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:94cf4285b46bc2d74fbe86e39a93ecfb3b0e584477fff7e38d293d47b88576e7", size = 7046132, upload-time = "2026-02-13T01:02:31.606Z" }, - { url = "https://files.pythonhosted.org/packages/b8/7b/c1612ec68d98e5f3dbc5b8a21ff5d40ab52409fcc89ea7afc8a197983297/mujoco-3.5.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:12bfb2bb70f760e0d51fd59f3c43b2906c7660a23954fd717321da52ba85a617", size = 6677917, upload-time = "2026-02-13T01:02:34.13Z" }, - { url = "https://files.pythonhosted.org/packages/c8/8a/229e4db3692be55532e155e2ca6a1363752243ee79df0e7e22ba00f716cf/mujoco-3.5.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:66fe37276644c28fab497929c55580725de81afc6d511a40cc27525a8dd99efa", size = 7170882, upload-time = "2026-02-13T01:02:36.086Z" }, - { url = "https://files.pythonhosted.org/packages/02/37/527d83610b878f27c01dd762e0e41aaa62f095c607f0500ac7f724a2c7a5/mujoco-3.5.0-cp312-cp312-win_amd64.whl", hash = "sha256:4b3a62af174ab59b9b6d816dca0786b7fd85ac081d6c2a931a2b22dd6e821f50", size = 5721886, upload-time = "2026-02-13T01:02:39.544Z" }, + { url = "https://files.pythonhosted.org/packages/38/c4/f8959e3d5d98b282e081ce08d07cd71ae949cc0ad9f2c39c0a69fcb88c8c/mujoco-3.6.0-cp312-cp312-macosx_10_16_x86_64.whl", hash = "sha256:e7e60ee4c07f6fecd63c23e6f47b8d7cdacad75d311739d50d50b5107a630af2", size = 7159624, upload-time = "2026-03-11T01:45:16.893Z" }, + { url = "https://files.pythonhosted.org/packages/26/55/7407eced2c44fbea233302d2c11e778852ea0f2eb0e14610f13a7e0d6ac7/mujoco-3.6.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ea71750f8cbe24b02a091093592f08fb71c95692b43c25e87dabe496ace0bb55", size = 7093719, upload-time = "2026-03-11T01:45:19.191Z" }, + { url = "https://files.pythonhosted.org/packages/2c/cc/2aae89c3a83fed29ccb9057c05fb4a218b2a42c6dea136d9a78fea6b39f8/mujoco-3.6.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:094de585a2084508f1cfd76170b0dfe1d9c122b3bd4677e96ef2383100c9032f", size = 6982824, upload-time = "2026-03-11T01:45:22.078Z" }, + { url = "https://files.pythonhosted.org/packages/52/6c/5ec4e93676a65064a6591176772e00cfa02716156a1d0a7d646a8203348f/mujoco-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8714fab312c7ee58f45bda7ef8762da2184e3a6a1d780a5093e93a160d66bd3d", size = 7473873, upload-time = "2026-03-11T01:45:24.671Z" }, + { url = "https://files.pythonhosted.org/packages/92/22/38d82f0c34213af53afbbb248b3442943ef48ffbac1e4c909b321e02ac56/mujoco-3.6.0-cp312-cp312-win_amd64.whl", hash = "sha256:3d4ec53e4e20fcc85843d607fa1648e0b12d2d2de81ee6f85926e95a7e84e8d8", size = 5764289, upload-time = "2026-03-11T01:45:27.014Z" }, ] [[package]] name = "mujoco-mjx" -version = "3.5.0" +version = "3.6.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "absl-py" }, @@ -969,9 +1304,34 @@ dependencies = [ { name = "scipy" }, { name = "trimesh" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/6c/3c/fc471adb5c83bb657c3634cf37c8c5cb5bb37c204d02192a4ee215132d1e/mujoco_mjx-3.5.0.tar.gz", hash = "sha256:42bdf3e80c0c4dfcfc78af97034f836d5292742e450a43a0dd9d44ada1e4bdc0", size = 6907429, upload-time = "2026-02-13T01:04:23.208Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ea/b2/f3fbed34c34d12c41463f4b621e2ad8e5907eddb197e96f692f4101d644c/mujoco_mjx-3.6.0.tar.gz", hash = "sha256:7cad0c40ebe63f18718d4171fe81e2fb978c81037ba859ceee7781ff81f36079", size = 6921702, upload-time = "2026-03-11T01:46:21.801Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/dd/ec/ba408121d07200f4d588ae83033a99dcd197bba47e35e50165d260f2ef6c/mujoco_mjx-3.5.0-py3-none-any.whl", hash = "sha256:633aa801f84fa2becc17ea124d95ad3e34f59fdfaa3720b7ec18b427f3c5bf46", size = 6992318, upload-time = "2026-02-13T01:04:21.21Z" }, + { url = "https://files.pythonhosted.org/packages/c2/7c/ad82beb7c4c9186d9fbef4799109d799692d70276bb1b3ee18a0674170d8/mujoco_mjx-3.6.0-py3-none-any.whl", hash = "sha256:c81000af0653f162b76009f48c153e9e6d19bfa8febe851e12466c81cbb7336a", size = 7013366, upload-time = "2026-03-11T01:46:19.148Z" }, +] + +[[package]] +name = "mujoco-warp" +version = "3.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "etils", extra = ["epath"] }, + { name = "mujoco" }, + { name = "numpy" }, + { name = "warp-lang" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0e/de/b853418268e9777cad2792ee3a145c8397e3d4517d136499645847ffd7f2/mujoco_warp-3.6.0.tar.gz", hash = "sha256:3c4111a4e13dc61268ddac52593ac5032c05a7d80f0c5e3c98bf5881e32b5d06", size = 1887269, upload-time = "2026-03-11T01:11:44.343Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a3/b5/06c1e23c0cc4a06da268aa5f0fe05348d89c9ecbd3ecfb4c6d2b14ea23b2/mujoco_warp-3.6.0-py3-none-any.whl", hash = "sha256:371a405b186332cbfaa9630aabf35967405ef847cb7ea7c018ec67426a2ea160", size = 1965960, upload-time = "2026-03-11T01:11:42.527Z" }, +] + +[[package]] +name = "mypy-extensions" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a2/6e/371856a3fb9d31ca8dac321cda606860fa4548858c0cc45d9d1d4ca2628b/mypy_extensions-1.1.0.tar.gz", hash = "sha256:52e68efc3284861e772bbcd66823fde5ae21fd2fdb51c62a211403730b916558", size = 6343, upload-time = "2025-04-22T14:54:24.164Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" }, ] [[package]] @@ -984,63 +1344,72 @@ wheels = [ ] [[package]] -name = "numpy" -version = "2.4.2" +name = "networkx" +version = "3.6.1" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/57/fd/0005efbd0af48e55eb3c7208af93f2862d4b1a56cd78e84309a2d959208d/numpy-2.4.2.tar.gz", hash = "sha256:659a6107e31a83c4e33f763942275fd278b21d095094044eb35569e86a21ddae", size = 20723651, upload-time = "2026-01-31T23:13:10.135Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/51/6e/6f394c9c77668153e14d4da83bcc247beb5952f6ead7699a1a2992613bea/numpy-2.4.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:21982668592194c609de53ba4933a7471880ccbaadcc52352694a59ecc860b3a", size = 16667963, upload-time = "2026-01-31T23:10:52.147Z" }, - { url = "https://files.pythonhosted.org/packages/1f/f8/55483431f2b2fd015ae6ed4fe62288823ce908437ed49db5a03d15151678/numpy-2.4.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40397bda92382fcec844066efb11f13e1c9a3e2a8e8f318fb72ed8b6db9f60f1", size = 14693571, upload-time = "2026-01-31T23:10:54.789Z" }, - { url = "https://files.pythonhosted.org/packages/2f/20/18026832b1845cdc82248208dd929ca14c9d8f2bac391f67440707fff27c/numpy-2.4.2-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:b3a24467af63c67829bfaa61eecf18d5432d4f11992688537be59ecd6ad32f5e", size = 5203469, upload-time = "2026-01-31T23:10:57.343Z" }, - { url = "https://files.pythonhosted.org/packages/7d/33/2eb97c8a77daaba34eaa3fa7241a14ac5f51c46a6bd5911361b644c4a1e2/numpy-2.4.2-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:805cc8de9fd6e7a22da5aed858e0ab16be5a4db6c873dde1d7451c541553aa27", size = 6550820, upload-time = "2026-01-31T23:10:59.429Z" }, - { url = "https://files.pythonhosted.org/packages/b1/91/b97fdfd12dc75b02c44e26c6638241cc004d4079a0321a69c62f51470c4c/numpy-2.4.2-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6d82351358ffbcdcd7b686b90742a9b86632d6c1c051016484fa0b326a0a1548", size = 15663067, upload-time = "2026-01-31T23:11:01.291Z" }, - { url = "https://files.pythonhosted.org/packages/f5/c6/a18e59f3f0b8071cc85cbc8d80cd02d68aa9710170b2553a117203d46936/numpy-2.4.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9e35d3e0144137d9fdae62912e869136164534d64a169f86438bc9561b6ad49f", size = 16619782, upload-time = "2026-01-31T23:11:03.669Z" }, - { url = "https://files.pythonhosted.org/packages/b7/83/9751502164601a79e18847309f5ceec0b1446d7b6aa12305759b72cf98b2/numpy-2.4.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:adb6ed2ad29b9e15321d167d152ee909ec73395901b70936f029c3bc6d7f4460", size = 17013128, upload-time = "2026-01-31T23:11:05.913Z" }, - { url = "https://files.pythonhosted.org/packages/61/c4/c4066322256ec740acc1c8923a10047818691d2f8aec254798f3dd90f5f2/numpy-2.4.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:8906e71fd8afcb76580404e2a950caef2685df3d2a57fe82a86ac8d33cc007ba", size = 18345324, upload-time = "2026-01-31T23:11:08.248Z" }, - { url = "https://files.pythonhosted.org/packages/ab/af/6157aa6da728fa4525a755bfad486ae7e3f76d4c1864138003eb84328497/numpy-2.4.2-cp312-cp312-win32.whl", hash = "sha256:ec055f6dae239a6299cace477b479cca2fc125c5675482daf1dd886933a1076f", size = 5960282, upload-time = "2026-01-31T23:11:10.497Z" }, - { url = "https://files.pythonhosted.org/packages/92/0f/7ceaaeaacb40567071e94dbf2c9480c0ae453d5bb4f52bea3892c39dc83c/numpy-2.4.2-cp312-cp312-win_amd64.whl", hash = "sha256:209fae046e62d0ce6435fcfe3b1a10537e858249b3d9b05829e2a05218296a85", size = 12314210, upload-time = "2026-01-31T23:11:12.176Z" }, - { url = "https://files.pythonhosted.org/packages/2f/a3/56c5c604fae6dd40fa2ed3040d005fca97e91bd320d232ac9931d77ba13c/numpy-2.4.2-cp312-cp312-win_arm64.whl", hash = "sha256:fbde1b0c6e81d56f5dccd95dd4a711d9b95df1ae4009a60887e56b27e8d903fa", size = 10220171, upload-time = "2026-01-31T23:11:14.684Z" }, + { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, +] + +[[package]] +name = "nodeenv" +version = "1.10.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/24/bf/d1bda4f6168e0b2e9e5958945e01910052158313224ada5ce1fb2e1113b8/nodeenv-1.10.0.tar.gz", hash = "sha256:996c191ad80897d076bdfba80a41994c2b47c68e224c542b48feba42ba00f8bb", size = 55611, upload-time = "2025-12-20T14:08:54.006Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl", hash = "sha256:5bb13e3eed2923615535339b3c620e76779af4cb4c6a90deccc9e36b274d3827", size = 23438, upload-time = "2025-12-20T14:08:52.782Z" }, +] + +[[package]] +name = "numpy" +version = "2.4.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/10/8b/c265f4823726ab832de836cdd184d0986dcf94480f81e8739692a7ac7af2/numpy-2.4.3.tar.gz", hash = "sha256:483a201202b73495f00dbc83796c6ae63137a9bdade074f7648b3e32613412dd", size = 20727743, upload-time = "2026-03-09T07:58:53.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a9/ed/6388632536f9788cea23a3a1b629f25b43eaacd7d7377e5d6bc7b9deb69b/numpy-2.4.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:61b0cbabbb6126c8df63b9a3a0c4b1f44ebca5e12ff6997b80fcf267fb3150ef", size = 16669628, upload-time = "2026-03-09T07:56:24.252Z" }, + { url = "https://files.pythonhosted.org/packages/74/1b/ee2abfc68e1ce728b2958b6ba831d65c62e1b13ce3017c13943f8f9b5b2e/numpy-2.4.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:7395e69ff32526710748f92cd8c9849b361830968ea3e24a676f272653e8983e", size = 14696872, upload-time = "2026-03-09T07:56:26.991Z" }, + { url = "https://files.pythonhosted.org/packages/ba/d1/780400e915ff5638166f11ca9dc2c5815189f3d7cf6f8759a1685e586413/numpy-2.4.3-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:abdce0f71dcb4a00e4e77f3faf05e4616ceccfe72ccaa07f47ee79cda3b7b0f4", size = 5203489, upload-time = "2026-03-09T07:56:29.414Z" }, + { url = "https://files.pythonhosted.org/packages/0b/bb/baffa907e9da4cc34a6e556d6d90e032f6d7a75ea47968ea92b4858826c4/numpy-2.4.3-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:48da3a4ee1336454b07497ff7ec83903efa5505792c4e6d9bf83d99dc07a1e18", size = 6550814, upload-time = "2026-03-09T07:56:32.225Z" }, + { url = "https://files.pythonhosted.org/packages/7b/12/8c9f0c6c95f76aeb20fc4a699c33e9f827fa0d0f857747c73bb7b17af945/numpy-2.4.3-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:32e3bef222ad6b052280311d1d60db8e259e4947052c3ae7dd6817451fc8a4c5", size = 15666601, upload-time = "2026-03-09T07:56:34.461Z" }, + { url = "https://files.pythonhosted.org/packages/bd/79/cc665495e4d57d0aa6fbcc0aa57aa82671dfc78fbf95fe733ed86d98f52a/numpy-2.4.3-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e7dd01a46700b1967487141a66ac1a3cf0dd8ebf1f08db37d46389401512ca97", size = 16621358, upload-time = "2026-03-09T07:56:36.852Z" }, + { url = "https://files.pythonhosted.org/packages/a8/40/b4ecb7224af1065c3539f5ecfff879d090de09608ad1008f02c05c770cb3/numpy-2.4.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:76f0f283506c28b12bba319c0fab98217e9f9b54e6160e9c79e9f7348ba32e9c", size = 17016135, upload-time = "2026-03-09T07:56:39.337Z" }, + { url = "https://files.pythonhosted.org/packages/f7/b1/6a88e888052eed951afed7a142dcdf3b149a030ca59b4c71eef085858e43/numpy-2.4.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:737f630a337364665aba3b5a77e56a68cc42d350edd010c345d65a3efa3addcc", size = 18345816, upload-time = "2026-03-09T07:56:42.31Z" }, + { url = "https://files.pythonhosted.org/packages/f3/8f/103a60c5f8c3d7fc678c19cd7b2476110da689ccb80bc18050efbaeae183/numpy-2.4.3-cp312-cp312-win32.whl", hash = "sha256:26952e18d82a1dbbc2f008d402021baa8d6fc8e84347a2072a25e08b46d698b9", size = 5960132, upload-time = "2026-03-09T07:56:44.851Z" }, + { url = "https://files.pythonhosted.org/packages/d7/7c/f5ee1bf6ed888494978046a809df2882aad35d414b622893322df7286879/numpy-2.4.3-cp312-cp312-win_amd64.whl", hash = "sha256:65f3c2455188f09678355f5cae1f959a06b778bc66d535da07bf2ef20cd319d5", size = 12316144, upload-time = "2026-03-09T07:56:47.057Z" }, + { url = "https://files.pythonhosted.org/packages/71/46/8d1cb3f7a00f2fb6394140e7e6623696e54c6318a9d9691bb4904672cf42/numpy-2.4.3-cp312-cp312-win_arm64.whl", hash = "sha256:2abad5c7fef172b3377502bde47892439bae394a71bc329f31df0fd829b41a9e", size = 10220364, upload-time = "2026-03-09T07:56:49.849Z" }, ] [[package]] name = "nvidia-cublas" -version = "13.2.1.1" +version = "13.1.0.3" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9d/36/0124129e1378e9834e0cbe19781fbe0ffd5f870c2af6f01cdf17a9869c39/nvidia_cublas-13.2.1.1-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:8b4a4cd8b73772fde9ccaa1f3967eb001ae5fde8b1dc37f7442d072b64d6f5da", size = 502470979, upload-time = "2026-01-13T22:39:37.619Z" }, - { url = "https://files.pythonhosted.org/packages/e2/e7/39e43c0688f9788c88da0b91ea18125448c5f515104aadf65a70243f144f/nvidia_cublas-13.2.1.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8c13c93cf8be4480b4909905c96d2d31575b4af43fcd3af0e84af94762665e4f", size = 401085577, upload-time = "2026-01-13T22:40:18.702Z" }, -] - -[[package]] -name = "nvidia-cuda-cccl" -version = "13.1.115" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ea/e8/0b0295d3f384d970ebffc2fbfe922c7e35bf1bc68f8d13a042f44935a667/nvidia_cuda_cccl-13.1.115-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5d5d26df9e56af547a3699048b466d55dbab7557af030161de96abd67834512a", size = 3480303, upload-time = "2026-01-13T22:29:23.018Z" }, - { url = "https://files.pythonhosted.org/packages/06/bf/7c2a9c40d4064e8e0e9dc14f4de357a810e7739ec00cc4f9144d7fc6d7af/nvidia_cuda_cccl-13.1.115-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d4a9e0590cd34290cda27f402a67fbf42e3ed043bdd927c6dd52419c860f2f92", size = 3430952, upload-time = "2026-01-13T22:29:51.639Z" }, + { url = "https://files.pythonhosted.org/packages/e1/a5/fce49e2ae977e0ccc084e5adafceb4f0ac0c8333cb6863501618a7277f67/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:c86fc7f7ae36d7528288c5d88098edcb7b02c633d262e7ddbb86b0ad91be5df2", size = 542851226, upload-time = "2025-10-09T08:59:04.818Z" }, + { url = "https://files.pythonhosted.org/packages/e7/44/423ac00af4dd95a5aeb27207e2c0d9b7118702149bf4704c3ddb55bb7429/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:ee8722c1f0145ab246bccb9e452153b5e0515fd094c3678df50b2a0888b8b171", size = 423133236, upload-time = "2025-10-09T08:59:32.536Z" }, ] [[package]] name = "nvidia-cuda-crt" -version = "13.1.115" +version = "13.2.51" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/26/b7/79cd270e5dcfd2339d0fb03d99ee96a8903e085201d01b0aa416d51ad710/nvidia_cuda_crt-13.1.115-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:eefd71d80d296391b100ab61e44bfd924308e2685d382f33a24d5c59213b43da", size = 132617, upload-time = "2026-01-13T22:30:07.696Z" }, - { url = "https://files.pythonhosted.org/packages/0e/d8/55a04550975b3c1b4bee09c0c7cfec45b3df9bafc0bb704412dc244c79ff/nvidia_cuda_crt-13.1.115-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:53c9b03f804dc7e539dd11914ec3d3fd0849791d8cd6a73f0d4df984bb339916", size = 132619, upload-time = "2026-01-13T22:30:37.924Z" }, + { url = "https://files.pythonhosted.org/packages/3e/89/34094e3b5eb0b12204ff97f8e4ee6a8df7b4a3e4811cace542fe361fe77c/nvidia_cuda_crt-13.2.51-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e6698ddc5da548ef7501f663ea55d18627999fa782a7b975f4dcd3fe3b26ef45", size = 133297, upload-time = "2026-03-09T09:28:55.788Z" }, + { url = "https://files.pythonhosted.org/packages/c8/5a/24af4197e8496870857fb56d5b93f65919fe5103fa311b526ec15d77a96a/nvidia_cuda_crt-13.2.51-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f4cda277fbf1025ad291a5d3b4dc4f788056ae11921552cdbebcf0626db99ba9", size = 133298, upload-time = "2026-03-09T09:29:25.823Z" }, ] [[package]] name = "nvidia-cuda-cupti" -version = "13.1.115" +version = "13.0.85" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/15/fd/5f1afe675d621a63d1d8505750ca2d934b62614465f7631f02fb9c6ea6f7/nvidia_cuda_cupti-13.1.115-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:2b0ffd140b48ca45d5b26e4b4bf718f193b9054f2e0b71f77490a6ddecfaaf27", size = 10922300, upload-time = "2026-01-13T22:32:20.783Z" }, - { url = "https://files.pythonhosted.org/packages/69/57/ac6c9c041331cd60df934806007f83407bf616135ab3b3677f396c3cf8e5/nvidia_cuda_cupti-13.1.115-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:2c76d91807cf76fc0e4b08ad71c4dff0e8e9f678fc8dbbbc43ec55125a70b553", size = 11372238, upload-time = "2026-01-13T22:32:49.132Z" }, + { url = "https://files.pythonhosted.org/packages/2a/2a/80353b103fc20ce05ef51e928daed4b6015db4aaa9162ed0997090fe2250/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:796bd679890ee55fb14a94629b698b6db54bcfd833d391d5e94017dd9d7d3151", size = 10310827, upload-time = "2025-09-04T08:26:42.012Z" }, + { url = "https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8", size = 10715597, upload-time = "2025-09-04T08:26:51.312Z" }, ] [[package]] name = "nvidia-cuda-nvcc" -version = "13.1.115" +version = "13.2.51" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-cuda-crt", marker = "sys_platform == 'linux'" }, @@ -1048,26 +1417,26 @@ dependencies = [ { name = "nvidia-nvvm", marker = "sys_platform == 'linux'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/54/11/31a1141e63dcb64ddedd056f2c454a9503e91674890dd4913e4be0b515f6/nvidia_cuda_nvcc-13.1.115-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ae46f71ea6f377e7719c1084ae8ca28d4e88770ee4436478c366c15f1fcfdebe", size = 34717342, upload-time = "2026-01-13T22:34:02.133Z" }, - { url = "https://files.pythonhosted.org/packages/be/05/52918fd34dabbc00290a89557cb2ee47f05f86a364d179375561222b460d/nvidia_cuda_nvcc-13.1.115-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fe5c0604869dfbc837bdce9e1eef8e4f0e0d103c29e3875014ba3538c70efb74", size = 42261253, upload-time = "2026-01-13T22:34:20.113Z" }, + { url = "https://files.pythonhosted.org/packages/4d/d8/3d1d733db86c1f18359151b0be0171b04738f17f09f98658caf9e3b5299d/nvidia_cuda_nvcc-13.2.51-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:48e070550a1290d696f055fa78443831bce5452cd2800eb3ab83f89b22c3b6cf", size = 38713648, upload-time = "2026-03-09T09:35:12.217Z" }, + { url = "https://files.pythonhosted.org/packages/5a/79/0da17b5b200ede8f25554f8c227c2624e26fb143c36ba7724b812c7e46ce/nvidia_cuda_nvcc-13.2.51-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:18aea9976c8a0033cc61d45baf5649a5bd8647a45999ddd50b885814a6190442", size = 44040269, upload-time = "2026-03-09T09:35:31.786Z" }, ] [[package]] name = "nvidia-cuda-nvrtc" -version = "13.1.115" +version = "13.0.88" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ff/d8/6fcf0f32d133a7da92efb1e90844d9f7c104627066cc52b13f7f0b128b54/nvidia_cuda_nvrtc-13.1.115-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:d7cf1284ab82f379884decc8813d9d4a729bc96b1ec020a9cf80303f698d73c4", size = 46564545, upload-time = "2026-01-13T22:35:53.834Z" }, - { url = "https://files.pythonhosted.org/packages/e9/f8/8bab039cbdd87af53f2ca0ca9e93bd676e53393ab4ea43da4735854dc1ce/nvidia_cuda_nvrtc-13.1.115-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:dfbc5e3bb19db41e4a05280b7b0cb9cbb624699f57dab3798455f43345541f99", size = 44308134, upload-time = "2026-01-13T22:35:35.287Z" }, + { url = "https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575", size = 90215200, upload-time = "2025-09-04T08:28:44.204Z" }, + { url = "https://files.pythonhosted.org/packages/b7/dc/6bb80850e0b7edd6588d560758f17e0550893a1feaf436807d64d2da040f/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d27f20a0ca67a4bb34268a5e951033496c5b74870b868bacd046b1b8e0c3267b", size = 43015449, upload-time = "2025-09-04T08:28:20.239Z" }, ] [[package]] name = "nvidia-cuda-runtime" -version = "13.1.80" +version = "13.0.96" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/32/b7/8075a985c5a4a828d850f244266c67134f60b28d1106cd4abe6187bedcc3/nvidia_cuda_runtime-13.1.80-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5f1aa8fd779cae1a78aa8fdbaa5e4f245300baff120099d9b5ee84e3f4506a29", size = 2317343, upload-time = "2025-12-05T17:34:27.824Z" }, - { url = "https://files.pythonhosted.org/packages/fe/60/03858bc3954b3263eedcff7626712c656b6b5a0d7d25bcb1fc1ebee9d4f1/nvidia_cuda_runtime-13.1.80-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:dac25b52e69050e5a7a84957e6df470a85b185f8f648a99a1ab1dfce027576cf", size = 2297816, upload-time = "2025-12-05T17:34:48.703Z" }, + { url = "https://files.pythonhosted.org/packages/87/4f/17d7b9b8e285199c58ce28e31b5c5bbaa4d8271af06a89b6405258245de2/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ef9bcbe90493a2b9d810e43d249adb3d02e98dd30200d86607d8d02687c43f55", size = 2261060, upload-time = "2025-10-09T08:55:15.78Z" }, + { url = "https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548", size = 2243632, upload-time = "2025-10-09T08:55:36.117Z" }, ] [[package]] @@ -1084,19 +1453,37 @@ wheels = [ [[package]] name = "nvidia-cufft" -version = "12.1.0.78" +version = "12.0.0.61" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/5c/87/787185b85241dc3d1bdbd20f2d315425836127d25b312a7e65ee839530cf/nvidia_cufft-12.1.0.78-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:639f60e802ace75277178148bc14bc17c6d1ec1f133cd81557c3e06ab70bf951", size = 226093383, upload-time = "2026-01-13T22:40:51.441Z" }, - { url = "https://files.pythonhosted.org/packages/5d/98/db43880ac42210eb85916446c8f6e6daf51ec538951fed654d2dc9fb327d/nvidia_cufft-12.1.0.78-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e1ceed544d2c1d92a19e91c3de20d7eaca58075e43d4b09efbda7f603607f1c2", size = 226114436, upload-time = "2026-01-13T22:41:31.803Z" }, + { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, + { url = "https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6c44f692dce8fd5ffd3e3df134b6cdb9c2f72d99cf40b62c32dde45eea9ddad3", size = 214085489, upload-time = "2025-09-04T08:31:56.044Z" }, +] + +[[package]] +name = "nvidia-cufile" +version = "1.15.1.6" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08a3ecefae5a01c7f5117351c64f17c7c62efa5fffdbe24fc7d298da19cd0b44", size = 1223672, upload-time = "2025-09-04T08:32:22.779Z" }, + { url = "https://files.pythonhosted.org/packages/ab/73/cc4a14c9813a8a0d509417cf5f4bdaba76e924d58beb9864f5a7baceefbf/nvidia_cufile-1.15.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:bdc0deedc61f548bddf7733bdc216456c2fdb101d020e1ab4b88d232d5e2f6d1", size = 1136992, upload-time = "2025-09-04T08:32:14.119Z" }, +] + +[[package]] +name = "nvidia-curand" +version = "10.4.0.35" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/72/7c2ae24fb6b63a32e6ae5d241cc65263ea18d08802aaae087d9f013335a2/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:133df5a7509c3e292aaa2b477afd0194f06ce4ea24d714d616ff36439cee349a", size = 61962106, upload-time = "2025-08-04T10:21:41.128Z" }, + { url = "https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:1aee33a5da6e1db083fe2b90082def8915f30f3248d5896bcec36a579d941bfc", size = 59544258, upload-time = "2025-08-04T10:22:03.992Z" }, ] [[package]] name = "nvidia-cusolver" -version = "12.0.9.81" +version = "12.0.4.66" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-cublas", marker = "sys_platform == 'linux'" }, @@ -1104,60 +1491,88 @@ dependencies = [ { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/57/01/b897eca0b6ba57eaf2ecb9f743549a06a6d520dae8c603262d1c164cae81/nvidia_cusolver-12.0.9.81-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:124e9ade619124f4a4a946dd63814d0cfc7b4738d9f34639b83f4d7c323e1eb9", size = 224479825, upload-time = "2026-01-13T22:43:45.06Z" }, - { url = "https://files.pythonhosted.org/packages/5b/75/3fb6ba779d87e9a507e8f3b5a64e2e6f6c83d45a86d8c9e67601cffcd117/nvidia_cusolver-12.0.9.81-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:5d28371e1301d1edfb15b8b1201e24cd680bc72de7f384cc162c7edbd8c5348c", size = 202032975, upload-time = "2026-01-13T22:44:23.405Z" }, + { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, + { url = "https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0a759da5dea5c0ea10fd307de75cdeb59e7ea4fcb8add0924859b944babf1112", size = 200941980, upload-time = "2025-09-04T08:33:22.767Z" }, ] [[package]] name = "nvidia-cusparse" -version = "12.7.3.1" +version = "12.6.3.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/98/b8/fbd4b324799fd8afb79642d5374ba7ab3eee0927257607f6b0631754f5bd/nvidia_cusparse-12.7.3.1-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c2cbabfe12fb58bfbd24fe2aa1b85a945900769eafc03cd788c5e19481f2a751", size = 170069494, upload-time = "2026-01-13T22:45:01.012Z" }, - { url = "https://files.pythonhosted.org/packages/4f/0c/6f1afab1aeba7c249f7ea2598513aaf22c30dde26e1c0c9681d9517b8ff3/nvidia_cusparse-12.7.3.1-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:39dc7926b835761ade1a7210e047877977cdf4153ac2f8df57662047916d92ad", size = 153162100, upload-time = "2026-01-13T22:45:38.173Z" }, + { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, + { url = "https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2b3c89c88d01ee0e477cb7f82ef60a11a4bcd57b6b87c33f789350b59759360b", size = 145942937, upload-time = "2025-09-04T08:33:58.029Z" }, +] + +[[package]] +name = "nvidia-cusparselt-cu13" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/10/8dcd1175260706a2fc92a16a52e306b71d4c1ea0b0cc4a9484183399818a/nvidia_cusparselt_cu13-0.8.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:400c6ed1cf6780fc6efedd64ec9f1345871767e6a1a0a552a1ea0578117ea77c", size = 220791277, upload-time = "2025-08-13T19:22:40.982Z" }, + { url = "https://files.pythonhosted.org/packages/fd/53/43b0d71f4e702fa9733f8b4571fdca50a8813f1e450b656c239beff12315/nvidia_cusparselt_cu13-0.8.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:25e30a8a7323935d4ad0340b95a0b69926eee755767e8e0b1cf8dd85b197d3fd", size = 169884119, upload-time = "2025-08-13T19:23:41.967Z" }, ] [[package]] name = "nvidia-nccl-cu13" -version = "2.29.3" +version = "2.28.9" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/27/59/ff243ebe6fa1767a9135719829347f609a90607cfbba9637ba3e9b3e36ce/nvidia_nccl_cu13-2.29.3-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:eab9f5c565ab3326906f1d1b5be5773a174c2a1b47002faed76f9e957392f713", size = 201042594, upload-time = "2026-02-03T21:10:54.736Z" }, - { url = "https://files.pythonhosted.org/packages/7b/70/aae7806eeaed043b3e212da435880ad067b5f14052986a6b4c0a4c62f68a/nvidia_nccl_cu13-2.29.3-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:2a321629f49490e4e0122ecb578a4b4a6f89e72740dd988e04dfa4758fab7fc3", size = 201104023, upload-time = "2026-02-03T21:11:24.071Z" }, + { url = "https://files.pythonhosted.org/packages/39/55/1920646a2e43ffd4fc958536b276197ed740e9e0c54105b4bb3521591fc7/nvidia_nccl_cu13-2.28.9-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:01c873ba1626b54caa12272ed228dc5b2781545e0ae8ba3f432a8ef1c6d78643", size = 196561677, upload-time = "2025-11-18T05:49:03.45Z" }, + { url = "https://files.pythonhosted.org/packages/b0/b4/878fefaad5b2bcc6fcf8d474a25e3e3774bc5133e4b58adff4d0bca238bc/nvidia_nccl_cu13-2.28.9-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:e4553a30f34195f3fa1da02a6da3d6337d28f2003943aa0a3d247bbc25fefc42", size = 196493177, upload-time = "2025-11-18T05:49:17.677Z" }, ] [[package]] name = "nvidia-nvjitlink" -version = "13.1.115" +version = "13.0.88" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/14/b0/1e00dd34da707dddaea23d8b367533e342ec84ce87a95c255bc8bee22121/nvidia_nvjitlink-13.1.115-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:658e87d42ac6be82ae2c56eb9abe45bc8de45b7cf7692d24f86848f0266d121b", size = 40920528, upload-time = "2026-01-13T22:48:29.661Z" }, - { url = "https://files.pythonhosted.org/packages/7c/51/8a946d47b8964e9cf1bab96cfa09b78f818a3b83231224509cd6b7a575b3/nvidia_nvjitlink-13.1.115-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:72bb753de9c968f2b4d885ac3e0129995f8701c48d24325bbf21b7a631d2d5c9", size = 38848076, upload-time = "2026-01-13T22:47:57.619Z" }, + { url = "https://files.pythonhosted.org/packages/56/7a/123e033aaff487c77107195fa5a2b8686795ca537935a24efae476c41f05/nvidia_nvjitlink-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:13a74f429e23b921c1109976abefacc69835f2f433ebd323d3946e11d804e47b", size = 40713933, upload-time = "2025-09-04T08:35:43.553Z" }, + { url = "https://files.pythonhosted.org/packages/ab/2c/93c5250e64df4f894f1cbb397c6fd71f79813f9fd79d7cd61de3f97b3c2d/nvidia_nvjitlink-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e931536ccc7d467a98ba1d8b89ff7fa7f1fa3b13f2b0069118cd7f47bff07d0c", size = 38768748, upload-time = "2025-09-04T08:35:20.008Z" }, ] [[package]] name = "nvidia-nvshmem-cu13" -version = "3.5.19" +version = "3.4.5" source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-cuda-cccl", marker = "sys_platform == 'linux'" }, -] wheels = [ - { url = "https://files.pythonhosted.org/packages/4e/62/d0835400d6d93474ed0b4655411eed2e807b8a1a8bb1e3a0e128e5209a63/nvidia_nvshmem_cu13-3.5.19-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:53440ae98eadd3bd83fc875c793d737e2ce618632759942a32d4cfb1814c0a74", size = 72222179, upload-time = "2026-01-02T04:24:34.97Z" }, - { url = "https://files.pythonhosted.org/packages/c8/0f/28e11b20c47e8b67d108c6895f4e3eba05071df16bc5bd2e11a572918330/nvidia_nvshmem_cu13-3.5.19-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0f734285d3a04824c13e9e2c1169dbb818be9cb0819de69389aaa067b8d18c93", size = 72433982, upload-time = "2026-01-02T04:25:05.69Z" }, + { url = "https://files.pythonhosted.org/packages/dc/0f/05cc9c720236dcd2db9c1ab97fff629e96821be2e63103569da0c9b72f19/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dc2a197f38e5d0376ad52cd1a2a3617d3cdc150fd5966f4aee9bcebb1d68fe9", size = 60215947, upload-time = "2025-09-06T00:32:20.022Z" }, + { url = "https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:290f0a2ee94c9f3687a02502f3b9299a9f9fe826e6d0287ee18482e78d495b80", size = 60412546, upload-time = "2025-09-06T00:32:41.564Z" }, +] + +[[package]] +name = "nvidia-nvtx" +version = "13.0.85" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4", size = 148047, upload-time = "2025-09-04T08:29:01.761Z" }, + { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" }, ] [[package]] name = "nvidia-nvvm" -version = "13.1.115" +version = "13.2.51" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3d/24/8ce3e8564e1ebf468b2b9a838f2a6c8f357ad63bec353a6a5c38b3f79250/nvidia_nvvm-13.1.115-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:6fdf483d85d9903f9cbf84c41d99221a86e0c0511d05de3d237b4ab69c11627e", size = 63386729, upload-time = "2026-01-13T22:51:00.152Z" }, - { url = "https://files.pythonhosted.org/packages/e6/01/7dc8dc16b86dd173dc786f454eff688efa4fd61e9eff5f582aad90621e4c/nvidia_nvvm-13.1.115-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:860706eb4b5851ac1d2bacf574b582ec66901acc31051da9b0cf9a2c910a9aa9", size = 61001803, upload-time = "2026-01-13T22:50:40.768Z" }, - { url = "https://files.pythonhosted.org/packages/bc/53/dc1cc7970beb56ec6ce3717fd110b04f57ee338556d3e06114a0bfbc267d/nvidia_nvvm-13.1.115-py3-none-win_amd64.whl", hash = "sha256:795c08d8a9d2f184dbd6ffcf049ce2ede5af1b77e7965e7936ef73125e7bee21", size = 55415346, upload-time = "2026-01-13T23:03:18.422Z" }, + { url = "https://files.pythonhosted.org/packages/34/4c/865325b6cffe2c2c20fe63696dca29b869ea7c0845aa743c217c2fb987dd/nvidia_nvvm-13.2.51-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:9c5725d97b1108bdb6c474784f7901c34f570319a2c2a0f279d23190070915f3", size = 64279456, upload-time = "2026-03-09T09:58:39.231Z" }, + { url = "https://files.pythonhosted.org/packages/8d/f2/c67ff35faf322d29a41046af76b4d9b86d8ac3f555f59d1a1defb7a4eca4/nvidia_nvvm-13.2.51-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:bcfd4be51f011045520974bd93f467bc2d64b87f333ccbdec883a372a55aa8f6", size = 61886052, upload-time = "2026-03-09T09:58:05.734Z" }, + { url = "https://files.pythonhosted.org/packages/ec/11/3f1ee9dce24b41812dd572a037c4436d4d21f759fbe373cc271b0ce98805/nvidia_nvvm-13.2.51-py3-none-win_amd64.whl", hash = "sha256:a4809baaa5429eabe1878853761ce31f0ba15216e2348710b7898dc591f5fc14", size = 56751075, upload-time = "2026-03-09T10:11:09.994Z" }, +] + +[[package]] +name = "omegaconf" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/09/48/6388f1bb9da707110532cb70ec4d2822858ddfb44f1cdf1233c20a80ea4b/omegaconf-2.3.0.tar.gz", hash = "sha256:d5d4b6d29955cc50ad50c46dc269bcd92c6e00f5f90d23ab5fee7bfca4ba4cc7", size = 3298120, upload-time = "2022-12-08T20:59:22.753Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e3/94/1843518e420fa3ed6919835845df698c7e27e183cb997394e4a670973a65/omegaconf-2.3.0-py3-none-any.whl", hash = "sha256:7b4df175cdb08ba400f45cae3bdcae7ba8365db4d165fc65fd04b050ab63b46b", size = 79500, upload-time = "2022-12-08T20:59:19.686Z" }, ] [[package]] @@ -1189,7 +1604,7 @@ wheels = [ [[package]] name = "optax" -version = "0.2.7" +version = "0.2.8" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "absl-py" }, @@ -1197,9 +1612,9 @@ dependencies = [ { name = "jaxlib" }, { name = "numpy" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/de/f7/a63fc3d262d7a58d7d53050dea1408a63738739569af34f8f754cf181ab1/optax-0.2.7.tar.gz", hash = "sha256:8b6b2e5bd62bcc6c11f6172a1aff0d86da0eaeecbd5465b2b366b5d3d64f6efc", size = 297524, upload-time = "2026-02-05T20:49:28.749Z" } +sdist = { url = "https://files.pythonhosted.org/packages/8c/f9/e3d11ae6f298ee941a0690e353a323d158ba5dedc436e75621c310845c5c/optax-0.2.8.tar.gz", hash = "sha256:5b225b35066fc3eebaa4d798f1b4173b4d57d1a480610908981f8343b50af0b0", size = 301193, upload-time = "2026-03-20T23:30:05.465Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b8/1e/94ad43e06887244b4d25f58b689122270ba3c129d3448052958eecf7518a/optax-0.2.7-py3-none-any.whl", hash = "sha256:241f2dfa104eab4fec2e16e7919f88df24a3da1481f95e264b3db396b30d4ff6", size = 399395, upload-time = "2026-02-05T20:49:26.883Z" }, + { url = "https://files.pythonhosted.org/packages/8a/69/6a93d8600c339d7687a05857c7907bd4dd8cf88691a5ea106d7a50af90a1/optax-0.2.8-py3-none-any.whl", hash = "sha256:e3ca2d36c99daab1800ae9dbc0545034382d6bc780b24d969e1b0df65fa31cb4", size = 402960, upload-time = "2026-03-20T23:30:03.886Z" }, ] [[package]] @@ -1227,6 +1642,26 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3a/85/41280ea5d6aa58d8033b2ac6ef70849dcbe37910b34b52c6195efb06ef9e/orbax_checkpoint-0.11.33-py3-none-any.whl", hash = "sha256:b8b6c40fe307d55c490c37852fcdc7ed86435613f40ff3887298454f667b58f1", size = 696815, upload-time = "2026-02-18T04:22:28.935Z" }, ] +[[package]] +name = "orbax-export" +version = "0.0.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "dataclasses-json" }, + { name = "etils" }, + { name = "jax" }, + { name = "jaxlib" }, + { name = "jaxtyping" }, + { name = "numpy" }, + { name = "orbax-checkpoint" }, + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1c/c8/ed7ac3c3c687bf129d7469b016c2b3d8777379f4ea453474e50ee41ce5cb/orbax_export-0.0.8.tar.gz", hash = "sha256:544eef564e2a6f17cd11b1167febe348b7b7cf56d9575de994a33d5613dd568a", size = 124980, upload-time = "2025-09-17T15:41:14.264Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/a9/3a755a58c8b6a36fe7e9e66bb6b93967ff49cdbc77cca8eacb2cf66435e9/orbax_export-0.0.8-py3-none-any.whl", hash = "sha256:f8037e1666ad28411cdb08d0668a2737b1281a32902c623ceda12109a089bc36", size = 180487, upload-time = "2025-09-17T15:41:12.928Z" }, +] + [[package]] name = "packaging" version = "26.0" @@ -1236,6 +1671,27 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/b9/c538f279a4e237a006a2c98387d081e9eb060d203d8ed34467cc0f0b9b53/packaging-26.0-py3-none-any.whl", hash = "sha256:b36f1fef9334a5588b4166f8bcd26a14e521f2b55e6b9de3aaa80d3ff7a37529", size = 74366, upload-time = "2026-01-21T20:50:37.788Z" }, ] +[[package]] +name = "pandas" +version = "3.0.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "python-dateutil" }, + { name = "tzdata", marker = "sys_platform == 'emscripten' or sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2e/0c/b28ed414f080ee0ad153f848586d61d1878f91689950f037f976ce15f6c8/pandas-3.0.1.tar.gz", hash = "sha256:4186a699674af418f655dbd420ed87f50d56b4cd6603784279d9eef6627823c8", size = 4641901, upload-time = "2026-02-17T22:20:16.434Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/37/51/b467209c08dae2c624873d7491ea47d2b47336e5403309d433ea79c38571/pandas-3.0.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:476f84f8c20c9f5bc47252b66b4bb25e1a9fc2fa98cead96744d8116cb85771d", size = 10344357, upload-time = "2026-02-17T22:18:38.262Z" }, + { url = "https://files.pythonhosted.org/packages/7c/f1/e2567ffc8951ab371db2e40b2fe068e36b81d8cf3260f06ae508700e5504/pandas-3.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0ab749dfba921edf641d4036c4c21c0b3ea70fea478165cb98a998fb2a261955", size = 9884543, upload-time = "2026-02-17T22:18:41.476Z" }, + { url = "https://files.pythonhosted.org/packages/d7/39/327802e0b6d693182403c144edacbc27eb82907b57062f23ef5a4c4a5ea7/pandas-3.0.1-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b8e36891080b87823aff3640c78649b91b8ff6eea3c0d70aeabd72ea43ab069b", size = 10396030, upload-time = "2026-02-17T22:18:43.822Z" }, + { url = "https://files.pythonhosted.org/packages/3d/fe/89d77e424365280b79d99b3e1e7d606f5165af2f2ecfaf0c6d24c799d607/pandas-3.0.1-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:532527a701281b9dd371e2f582ed9094f4c12dd9ffb82c0c54ee28d8ac9520c4", size = 10876435, upload-time = "2026-02-17T22:18:45.954Z" }, + { url = "https://files.pythonhosted.org/packages/b5/a6/2a75320849dd154a793f69c951db759aedb8d1dd3939eeacda9bdcfa1629/pandas-3.0.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:356e5c055ed9b0da1580d465657bc7d00635af4fd47f30afb23025352ba764d1", size = 11405133, upload-time = "2026-02-17T22:18:48.533Z" }, + { url = "https://files.pythonhosted.org/packages/58/53/1d68fafb2e02d7881df66aa53be4cd748d25cbe311f3b3c85c93ea5d30ca/pandas-3.0.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:9d810036895f9ad6345b8f2a338dd6998a74e8483847403582cab67745bff821", size = 11932065, upload-time = "2026-02-17T22:18:50.837Z" }, + { url = "https://files.pythonhosted.org/packages/75/08/67cc404b3a966b6df27b38370ddd96b3b023030b572283d035181854aac5/pandas-3.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:536232a5fe26dd989bd633e7a0c450705fdc86a207fec7254a55e9a22950fe43", size = 9741627, upload-time = "2026-02-17T22:18:53.905Z" }, + { url = "https://files.pythonhosted.org/packages/86/4f/caf9952948fb00d23795f09b893d11f1cacb384e666854d87249530f7cbe/pandas-3.0.1-cp312-cp312-win_arm64.whl", hash = "sha256:0f463ebfd8de7f326d38037c7363c6dacb857c5881ab8961fb387804d6daf2f7", size = 9052483, upload-time = "2026-02-17T22:18:57.31Z" }, +] + [[package]] name = "parso" version = "0.8.6" @@ -1250,7 +1706,7 @@ name = "pexpect" version = "4.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "ptyprocess" }, + { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } wheels = [ @@ -1278,11 +1734,36 @@ wheels = [ [[package]] name = "platformdirs" -version = "4.9.2" +version = "4.9.4" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/1b/04/fea538adf7dbbd6d186f551d595961e564a3b6715bdf276b477460858672/platformdirs-4.9.2.tar.gz", hash = "sha256:9a33809944b9db043ad67ca0db94b14bf452cc6aeaac46a88ea55b26e2e9d291", size = 28394, upload-time = "2026-02-16T03:56:10.574Z" } +sdist = { url = "https://files.pythonhosted.org/packages/19/56/8d4c30c8a1d07013911a8fdbd8f89440ef9f08d07a1b50ab8ca8be5a20f9/platformdirs-4.9.4.tar.gz", hash = "sha256:1ec356301b7dc906d83f371c8f487070e99d3ccf9e501686456394622a01a934", size = 28737, upload-time = "2026-03-05T18:34:13.271Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/48/31/05e764397056194206169869b50cf2fee4dbbbc71b344705b9c0d878d4d8/platformdirs-4.9.2-py3-none-any.whl", hash = "sha256:9170634f126f8efdae22fb58ae8a0eaa86f38365bc57897a6c4f781d1f5875bd", size = 21168, upload-time = "2026-02-16T03:56:08.891Z" }, + { url = "https://files.pythonhosted.org/packages/63/d7/97f7e3a6abb67d8080dd406fd4df842c2be0efaf712d1c899c32a075027c/platformdirs-4.9.4-py3-none-any.whl", hash = "sha256:68a9a4619a666ea6439f2ff250c12a853cd1cbd5158d258bd824a7df6be2f868", size = 21216, upload-time = "2026-03-05T18:34:12.172Z" }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + +[[package]] +name = "pre-commit" +version = "4.5.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cfgv" }, + { name = "identify" }, + { name = "nodeenv" }, + { name = "pyyaml" }, + { name = "virtualenv" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/40/f1/6d86a29246dfd2e9b6237f0b5823717f60cad94d47ddc26afa916d21f525/pre_commit-4.5.1.tar.gz", hash = "sha256:eb545fcff725875197837263e977ea257a402056661f09dae08e4b149b030a61", size = 198232, upload-time = "2025-12-16T21:14:33.552Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5d/19/fd3ef348460c80af7bb4669ea7926651d1f95c23ff2df18b9d24bab4f3fa/pre_commit-4.5.1-py2.py3-none-any.whl", hash = "sha256:3b3afd891e97337708c1674210f8eba659b52a38ea5f822ff142d10786221f77", size = 226437, upload-time = "2025-12-16T21:14:32.409Z" }, ] [[package]] @@ -1299,17 +1780,17 @@ wheels = [ [[package]] name = "protobuf" -version = "6.33.5" +version = "6.33.6" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ba/25/7c72c307aafc96fa87062aa6291d9f7c94836e43214d43722e86037aac02/protobuf-6.33.5.tar.gz", hash = "sha256:6ddcac2a081f8b7b9642c09406bc6a4290128fce5f471cddd165960bb9119e5c", size = 444465, upload-time = "2026-01-29T21:51:33.494Z" } +sdist = { url = "https://files.pythonhosted.org/packages/66/70/e908e9c5e52ef7c3a6c7902c9dfbb34c7e29c25d2f81ade3856445fd5c94/protobuf-6.33.6.tar.gz", hash = "sha256:a6768d25248312c297558af96a9f9c929e8c4cee0659cb07e780731095f38135", size = 444531, upload-time = "2026-03-18T19:05:00.988Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b1/79/af92d0a8369732b027e6d6084251dd8e782c685c72da161bd4a2e00fbabb/protobuf-6.33.5-cp310-abi3-win32.whl", hash = "sha256:d71b040839446bac0f4d162e758bea99c8251161dae9d0983a3b88dee345153b", size = 425769, upload-time = "2026-01-29T21:51:21.751Z" }, - { url = "https://files.pythonhosted.org/packages/55/75/bb9bc917d10e9ee13dee8607eb9ab963b7cf8be607c46e7862c748aa2af7/protobuf-6.33.5-cp310-abi3-win_amd64.whl", hash = "sha256:3093804752167bcab3998bec9f1048baae6e29505adaf1afd14a37bddede533c", size = 437118, upload-time = "2026-01-29T21:51:24.022Z" }, - { url = "https://files.pythonhosted.org/packages/a2/6b/e48dfc1191bc5b52950246275bf4089773e91cb5ba3592621723cdddca62/protobuf-6.33.5-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:a5cb85982d95d906df1e2210e58f8e4f1e3cdc088e52c921a041f9c9a0386de5", size = 427766, upload-time = "2026-01-29T21:51:25.413Z" }, - { url = "https://files.pythonhosted.org/packages/4e/b1/c79468184310de09d75095ed1314b839eb2f72df71097db9d1404a1b2717/protobuf-6.33.5-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:9b71e0281f36f179d00cbcb119cb19dec4d14a81393e5ea220f64b286173e190", size = 324638, upload-time = "2026-01-29T21:51:26.423Z" }, - { url = "https://files.pythonhosted.org/packages/c5/f5/65d838092fd01c44d16037953fd4c2cc851e783de9b8f02b27ec4ffd906f/protobuf-6.33.5-cp39-abi3-manylinux2014_s390x.whl", hash = "sha256:8afa18e1d6d20af15b417e728e9f60f3aa108ee76f23c3b2c07a2c3b546d3afd", size = 339411, upload-time = "2026-01-29T21:51:27.446Z" }, - { url = "https://files.pythonhosted.org/packages/9b/53/a9443aa3ca9ba8724fdfa02dd1887c1bcd8e89556b715cfbacca6b63dbec/protobuf-6.33.5-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:cbf16ba3350fb7b889fca858fb215967792dc125b35c7976ca4818bee3521cf0", size = 323465, upload-time = "2026-01-29T21:51:28.925Z" }, - { url = "https://files.pythonhosted.org/packages/57/bf/2086963c69bdac3d7cff1cc7ff79b8ce5ea0bec6797a017e1be338a46248/protobuf-6.33.5-py3-none-any.whl", hash = "sha256:69915a973dd0f60f31a08b8318b73eab2bd6a392c79184b3612226b0a3f8ec02", size = 170687, upload-time = "2026-01-29T21:51:32.557Z" }, + { url = "https://files.pythonhosted.org/packages/fc/9f/2f509339e89cfa6f6a4c4ff50438db9ca488dec341f7e454adad60150b00/protobuf-6.33.6-cp310-abi3-win32.whl", hash = "sha256:7d29d9b65f8afef196f8334e80d6bc1d5d4adedb449971fefd3723824e6e77d3", size = 425739, upload-time = "2026-03-18T19:04:48.373Z" }, + { url = "https://files.pythonhosted.org/packages/76/5d/683efcd4798e0030c1bab27374fd13a89f7c2515fb1f3123efdfaa5eab57/protobuf-6.33.6-cp310-abi3-win_amd64.whl", hash = "sha256:0cd27b587afca21b7cfa59a74dcbd48a50f0a6400cfb59391340ad729d91d326", size = 437089, upload-time = "2026-03-18T19:04:50.381Z" }, + { url = "https://files.pythonhosted.org/packages/5c/01/a3c3ed5cd186f39e7880f8303cc51385a198a81469d53d0fdecf1f64d929/protobuf-6.33.6-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:9720e6961b251bde64edfdab7d500725a2af5280f3f4c87e57c0208376aa8c3a", size = 427737, upload-time = "2026-03-18T19:04:51.866Z" }, + { url = "https://files.pythonhosted.org/packages/ee/90/b3c01fdec7d2f627b3a6884243ba328c1217ed2d978def5c12dc50d328a3/protobuf-6.33.6-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:e2afbae9b8e1825e3529f88d514754e094278bb95eadc0e199751cdd9a2e82a2", size = 324610, upload-time = "2026-03-18T19:04:53.096Z" }, + { url = "https://files.pythonhosted.org/packages/9b/ca/25afc144934014700c52e05103c2421997482d561f3101ff352e1292fb81/protobuf-6.33.6-cp39-abi3-manylinux2014_s390x.whl", hash = "sha256:c96c37eec15086b79762ed265d59ab204dabc53056e3443e702d2681f4b39ce3", size = 339381, upload-time = "2026-03-18T19:04:54.616Z" }, + { url = "https://files.pythonhosted.org/packages/16/92/d1e32e3e0d894fe00b15ce28ad4944ab692713f2e7f0a99787405e43533a/protobuf-6.33.6-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:e9db7e292e0ab79dd108d7f1a94fe31601ce1ee3f7b79e0692043423020b0593", size = 323436, upload-time = "2026-03-18T19:04:55.768Z" }, + { url = "https://files.pythonhosted.org/packages/c4/72/02445137af02769918a93807b2b7890047c32bfb9f90371cbc12688819eb/protobuf-6.33.6-py3-none-any.whl", hash = "sha256:77179e006c476e69bf8e8ce866640091ec42e1beb80b213c3900006ecfba6901", size = 170656, upload-time = "2026-03-18T19:04:59.826Z" }, ] [[package]] @@ -1439,6 +1920,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/10/bd/c038d7cc38edc1aa5bf91ab8068b63d4308c66c4c8bb3cbba7dfbc049f9c/pyparsing-3.3.2-py3-none-any.whl", hash = "sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d", size = 122781, upload-time = "2026-01-21T03:57:55.912Z" }, ] +[[package]] +name = "pytest" +version = "9.0.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d1/db/7ef3487e0fb0049ddb5ce41d3a49c235bf9ad299b6a25d5780a89f19230f/pytest-9.0.2.tar.gz", hash = "sha256:75186651a92bd89611d1d9fc20f0b4345fd827c41ccd5c299a868a05d70edf11", size = 1568901, upload-time = "2025-12-06T21:30:51.014Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl", hash = "sha256:711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b", size = 374801, upload-time = "2025-12-06T21:30:49.154Z" }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0" @@ -1451,6 +1948,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, ] +[[package]] +name = "python-discovery" +version = "1.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "filelock" }, + { name = "platformdirs" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9c/90/bcce6b46823c9bec1757c964dc37ed332579be512e17a30e9698095dcae4/python_discovery-1.2.0.tar.gz", hash = "sha256:7d33e350704818b09e3da2bd419d37e21e7c30db6e0977bb438916e06b41b5b1", size = 58055, upload-time = "2026-03-19T01:43:08.248Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c2/3c/2005227cb951df502412de2fa781f800663cccbef8d90ec6f1b371ac2c0d/python_discovery-1.2.0-py3-none-any.whl", hash = "sha256:1e108f1bbe2ed0ef089823d28805d5ad32be8e734b86a5f212bf89b71c266e4a", size = 31524, upload-time = "2026-03-19T01:43:07.045Z" }, +] + [[package]] name = "pyyaml" version = "6.0.3" @@ -1520,27 +2030,27 @@ wheels = [ [[package]] name = "ruff" -version = "0.15.2" +version = "0.15.7" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/06/04/eab13a954e763b0606f460443fcbf6bb5a0faf06890ea3754ff16523dce5/ruff-0.15.2.tar.gz", hash = "sha256:14b965afee0969e68bb871eba625343b8673375f457af4abe98553e8bbb98342", size = 4558148, upload-time = "2026-02-19T22:32:20.271Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a1/22/9e4f66ee588588dc6c9af6a994e12d26e19efbe874d1a909d09a6dac7a59/ruff-0.15.7.tar.gz", hash = "sha256:04f1ae61fc20fe0b148617c324d9d009b5f63412c0b16474f3d5f1a1a665f7ac", size = 4601277, upload-time = "2026-03-19T16:26:22.605Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/2f/70/3a4dc6d09b13cb3e695f28307e5d889b2e1a66b7af9c5e257e796695b0e6/ruff-0.15.2-py3-none-linux_armv6l.whl", hash = "sha256:120691a6fdae2f16d65435648160f5b81a9625288f75544dc40637436b5d3c0d", size = 10430565, upload-time = "2026-02-19T22:32:41.824Z" }, - { url = "https://files.pythonhosted.org/packages/71/0b/bb8457b56185ece1305c666dc895832946d24055be90692381c31d57466d/ruff-0.15.2-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:a89056d831256099658b6bba4037ac6dd06f49d194199215befe2bb10457ea5e", size = 10820354, upload-time = "2026-02-19T22:32:07.366Z" }, - { url = "https://files.pythonhosted.org/packages/2d/c1/e0532d7f9c9e0b14c46f61b14afd563298b8b83f337b6789ddd987e46121/ruff-0.15.2-py3-none-macosx_11_0_arm64.whl", hash = "sha256:e36dee3a64be0ebd23c86ffa3aa3fd3ac9a712ff295e192243f814a830b6bd87", size = 10170767, upload-time = "2026-02-19T22:32:13.188Z" }, - { url = "https://files.pythonhosted.org/packages/47/e8/da1aa341d3af017a21c7a62fb5ec31d4e7ad0a93ab80e3a508316efbcb23/ruff-0.15.2-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a9fb47b6d9764677f8c0a193c0943ce9a05d6763523f132325af8a858eadc2b9", size = 10529591, upload-time = "2026-02-19T22:32:02.547Z" }, - { url = "https://files.pythonhosted.org/packages/93/74/184fbf38e9f3510231fbc5e437e808f0b48c42d1df9434b208821efcd8d6/ruff-0.15.2-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:f376990f9d0d6442ea9014b19621d8f2aaf2b8e39fdbfc79220b7f0c596c9b80", size = 10260771, upload-time = "2026-02-19T22:32:36.938Z" }, - { url = "https://files.pythonhosted.org/packages/05/ac/605c20b8e059a0bc4b42360414baa4892ff278cec1c91fff4be0dceedefd/ruff-0.15.2-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2dcc987551952d73cbf5c88d9fdee815618d497e4df86cd4c4824cc59d5dd75f", size = 11045791, upload-time = "2026-02-19T22:32:31.642Z" }, - { url = "https://files.pythonhosted.org/packages/fd/52/db6e419908f45a894924d410ac77d64bdd98ff86901d833364251bd08e22/ruff-0.15.2-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:42a47fd785cbe8c01b9ff45031af875d101b040ad8f4de7bbb716487c74c9a77", size = 11879271, upload-time = "2026-02-19T22:32:29.305Z" }, - { url = "https://files.pythonhosted.org/packages/3e/d8/7992b18f2008bdc9231d0f10b16df7dda964dbf639e2b8b4c1b4e91b83af/ruff-0.15.2-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:cbe9f49354866e575b4c6943856989f966421870e85cd2ac94dccb0a9dcb2fea", size = 11303707, upload-time = "2026-02-19T22:32:22.492Z" }, - { url = "https://files.pythonhosted.org/packages/d7/02/849b46184bcfdd4b64cde61752cc9a146c54759ed036edd11857e9b8443b/ruff-0.15.2-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b7a672c82b5f9887576087d97be5ce439f04bbaf548ee987b92d3a7dede41d3a", size = 11149151, upload-time = "2026-02-19T22:32:44.234Z" }, - { url = "https://files.pythonhosted.org/packages/70/04/f5284e388bab60d1d3b99614a5a9aeb03e0f333847e2429bebd2aaa1feec/ruff-0.15.2-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:72ecc64f46f7019e2bcc3cdc05d4a7da958b629a5ab7033195e11a438403d956", size = 11091132, upload-time = "2026-02-19T22:32:24.691Z" }, - { url = "https://files.pythonhosted.org/packages/fa/ae/88d844a21110e14d92cf73d57363fab59b727ebeabe78009b9ccb23500af/ruff-0.15.2-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:8dcf243b15b561c655c1ef2f2b0050e5d50db37fe90115507f6ff37d865dc8b4", size = 10504717, upload-time = "2026-02-19T22:32:26.75Z" }, - { url = "https://files.pythonhosted.org/packages/64/27/867076a6ada7f2b9c8292884ab44d08fd2ba71bd2b5364d4136f3cd537e1/ruff-0.15.2-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:dab6941c862c05739774677c6273166d2510d254dac0695c0e3f5efa1b5585de", size = 10263122, upload-time = "2026-02-19T22:32:10.036Z" }, - { url = "https://files.pythonhosted.org/packages/e7/ef/faf9321d550f8ebf0c6373696e70d1758e20ccdc3951ad7af00c0956be7c/ruff-0.15.2-py3-none-musllinux_1_2_i686.whl", hash = "sha256:1b9164f57fc36058e9a6806eb92af185b0697c9fe4c7c52caa431c6554521e5c", size = 10735295, upload-time = "2026-02-19T22:32:39.227Z" }, - { url = "https://files.pythonhosted.org/packages/2f/55/e8089fec62e050ba84d71b70e7834b97709ca9b7aba10c1a0b196e493f97/ruff-0.15.2-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:80d24fcae24d42659db7e335b9e1531697a7102c19185b8dc4a028b952865fd8", size = 11241641, upload-time = "2026-02-19T22:32:34.617Z" }, - { url = "https://files.pythonhosted.org/packages/23/01/1c30526460f4d23222d0fabd5888868262fd0e2b71a00570ca26483cd993/ruff-0.15.2-py3-none-win32.whl", hash = "sha256:fd5ff9e5f519a7e1bd99cbe8daa324010a74f5e2ebc97c6242c08f26f3714f6f", size = 10507885, upload-time = "2026-02-19T22:32:15.635Z" }, - { url = "https://files.pythonhosted.org/packages/5c/10/3d18e3bbdf8fc50bbb4ac3cc45970aa5a9753c5cb51bf9ed9a3cd8b79fa3/ruff-0.15.2-py3-none-win_amd64.whl", hash = "sha256:d20014e3dfa400f3ff84830dfb5755ece2de45ab62ecea4af6b7262d0fb4f7c5", size = 11623725, upload-time = "2026-02-19T22:32:04.947Z" }, - { url = "https://files.pythonhosted.org/packages/6d/78/097c0798b1dab9f8affe73da9642bb4500e098cb27fd8dc9724816ac747b/ruff-0.15.2-py3-none-win_arm64.whl", hash = "sha256:cabddc5822acdc8f7b5527b36ceac55cc51eec7b1946e60181de8fe83ca8876e", size = 10941649, upload-time = "2026-02-19T22:32:18.108Z" }, + { url = "https://files.pythonhosted.org/packages/41/2f/0b08ced94412af091807b6119ca03755d651d3d93a242682bf020189db94/ruff-0.15.7-py3-none-linux_armv6l.whl", hash = "sha256:a81cc5b6910fb7dfc7c32d20652e50fa05963f6e13ead3c5915c41ac5d16668e", size = 10489037, upload-time = "2026-03-19T16:26:32.47Z" }, + { url = "https://files.pythonhosted.org/packages/91/4a/82e0fa632e5c8b1eba5ee86ecd929e8ff327bbdbfb3c6ac5d81631bef605/ruff-0.15.7-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:722d165bd52403f3bdabc0ce9e41fc47070ac56d7a91b4e0d097b516a53a3477", size = 10955433, upload-time = "2026-03-19T16:27:00.205Z" }, + { url = "https://files.pythonhosted.org/packages/ab/10/12586735d0ff42526ad78c049bf51d7428618c8b5c467e72508c694119df/ruff-0.15.7-py3-none-macosx_11_0_arm64.whl", hash = "sha256:7fbc2448094262552146cbe1b9643a92f66559d3761f1ad0656d4991491af49e", size = 10269302, upload-time = "2026-03-19T16:26:26.183Z" }, + { url = "https://files.pythonhosted.org/packages/eb/5d/32b5c44ccf149a26623671df49cbfbd0a0ae511ff3df9d9d2426966a8d57/ruff-0.15.7-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6b39329b60eba44156d138275323cc726bbfbddcec3063da57caa8a8b1d50adf", size = 10607625, upload-time = "2026-03-19T16:27:03.263Z" }, + { url = "https://files.pythonhosted.org/packages/5d/f1/f0001cabe86173aaacb6eb9bb734aa0605f9a6aa6fa7d43cb49cbc4af9c9/ruff-0.15.7-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:87768c151808505f2bfc93ae44e5f9e7c8518943e5074f76ac21558ef5627c85", size = 10324743, upload-time = "2026-03-19T16:27:09.791Z" }, + { url = "https://files.pythonhosted.org/packages/7a/87/b8a8f3d56b8d848008559e7c9d8bf367934d5367f6d932ba779456e2f73b/ruff-0.15.7-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:fb0511670002c6c529ec66c0e30641c976c8963de26a113f3a30456b702468b0", size = 11138536, upload-time = "2026-03-19T16:27:06.101Z" }, + { url = "https://files.pythonhosted.org/packages/e4/f2/4fd0d05aab0c5934b2e1464784f85ba2eab9d54bffc53fb5430d1ed8b829/ruff-0.15.7-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:e0d19644f801849229db8345180a71bee5407b429dd217f853ec515e968a6912", size = 11994292, upload-time = "2026-03-19T16:26:48.718Z" }, + { url = "https://files.pythonhosted.org/packages/64/22/fc4483871e767e5e95d1622ad83dad5ebb830f762ed0420fde7dfa9d9b08/ruff-0.15.7-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:4806d8e09ef5e84eb19ba833d0442f7e300b23fe3f0981cae159a248a10f0036", size = 11398981, upload-time = "2026-03-19T16:26:54.513Z" }, + { url = "https://files.pythonhosted.org/packages/b0/99/66f0343176d5eab02c3f7fcd2de7a8e0dd7a41f0d982bee56cd1c24db62b/ruff-0.15.7-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dce0896488562f09a27b9c91b1f58a097457143931f3c4d519690dea54e624c5", size = 11242422, upload-time = "2026-03-19T16:26:29.277Z" }, + { url = "https://files.pythonhosted.org/packages/5d/3a/a7060f145bfdcce4c987ea27788b30c60e2c81d6e9a65157ca8afe646328/ruff-0.15.7-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:1852ce241d2bc89e5dc823e03cff4ce73d816b5c6cdadd27dbfe7b03217d2a12", size = 11232158, upload-time = "2026-03-19T16:26:42.321Z" }, + { url = "https://files.pythonhosted.org/packages/a7/53/90fbb9e08b29c048c403558d3cdd0adf2668b02ce9d50602452e187cd4af/ruff-0.15.7-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:5f3e4b221fb4bd293f79912fc5e93a9063ebd6d0dcbd528f91b89172a9b8436c", size = 10577861, upload-time = "2026-03-19T16:26:57.459Z" }, + { url = "https://files.pythonhosted.org/packages/2f/aa/5f486226538fe4d0f0439e2da1716e1acf895e2a232b26f2459c55f8ddad/ruff-0.15.7-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:b15e48602c9c1d9bdc504b472e90b90c97dc7d46c7028011ae67f3861ceba7b4", size = 10327310, upload-time = "2026-03-19T16:26:35.909Z" }, + { url = "https://files.pythonhosted.org/packages/99/9e/271afdffb81fe7bfc8c43ba079e9d96238f674380099457a74ccb3863857/ruff-0.15.7-py3-none-musllinux_1_2_i686.whl", hash = "sha256:1b4705e0e85cedc74b0a23cf6a179dbb3df184cb227761979cc76c0440b5ab0d", size = 10840752, upload-time = "2026-03-19T16:26:45.723Z" }, + { url = "https://files.pythonhosted.org/packages/bf/29/a4ae78394f76c7759953c47884eb44de271b03a66634148d9f7d11e721bd/ruff-0.15.7-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:112c1fa316a558bb34319282c1200a8bf0495f1b735aeb78bfcb2991e6087580", size = 11336961, upload-time = "2026-03-19T16:26:39.076Z" }, + { url = "https://files.pythonhosted.org/packages/26/6b/8786ba5736562220d588a2f6653e6c17e90c59ced34a2d7b512ef8956103/ruff-0.15.7-py3-none-win32.whl", hash = "sha256:6d39e2d3505b082323352f733599f28169d12e891f7dd407f2d4f54b4c2886de", size = 10582538, upload-time = "2026-03-19T16:26:15.992Z" }, + { url = "https://files.pythonhosted.org/packages/2b/e9/346d4d3fffc6871125e877dae8d9a1966b254fbd92a50f8561078b88b099/ruff-0.15.7-py3-none-win_amd64.whl", hash = "sha256:4d53d712ddebcd7dace1bc395367aec12c057aacfe9adbb6d832302575f4d3a1", size = 11755839, upload-time = "2026-03-19T16:26:19.897Z" }, + { url = "https://files.pythonhosted.org/packages/8f/e8/726643a3ea68c727da31570bde48c7a10f1aa60eddd628d94078fec586ff/ruff-0.15.7-py3-none-win_arm64.whl", hash = "sha256:18e8d73f1c3fdf27931497972250340f92e8c861722161a9caeb89a58ead6ed2", size = 11023304, upload-time = "2026-03-19T16:26:51.669Z" }, ] [[package]] @@ -1564,26 +2074,40 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bc/98/fe9ae9ffb3b54b62559f52dedaebe204b408db8109a8c66fdd04869e6424/scipy-1.17.1-cp312-cp312-win_arm64.whl", hash = "sha256:f4115102802df98b2b0db3cce5cb9b92572633a1197c77b7553e5203f284a5b3", size = 24547340, upload-time = "2026-02-23T00:19:12.024Z" }, ] +[[package]] +name = "seaborn" +version = "0.13.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "matplotlib" }, + { name = "numpy" }, + { name = "pandas" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/86/59/a451d7420a77ab0b98f7affa3a1d78a313d2f7281a57afb1a34bae8ab412/seaborn-0.13.2.tar.gz", hash = "sha256:93e60a40988f4d65e9f4885df477e2fdaff6b73a9ded434c1ab356dd57eefff7", size = 1457696, upload-time = "2024-01-25T13:21:52.551Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/83/11/00d3c3dfc25ad54e731d91449895a79e4bf2384dc3ac01809010ba88f6d5/seaborn-0.13.2-py3-none-any.whl", hash = "sha256:636f8336facf092165e27924f223d3c62ca560b1f2bb5dff7ab7fad265361987", size = 294914, upload-time = "2024-01-25T13:21:49.598Z" }, +] + [[package]] name = "sentry-sdk" -version = "2.53.0" +version = "2.56.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "certifi" }, { name = "urllib3" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d3/06/66c8b705179bc54087845f28fd1b72f83751b6e9a195628e2e9af9926505/sentry_sdk-2.53.0.tar.gz", hash = "sha256:6520ef2c4acd823f28efc55e43eb6ce2e6d9f954a95a3aa96b6fd14871e92b77", size = 412369, upload-time = "2026-02-16T11:11:14.743Z" } +sdist = { url = "https://files.pythonhosted.org/packages/de/df/5008954f5466085966468612a7d1638487596ee6d2fd7fb51783a85351bf/sentry_sdk-2.56.0.tar.gz", hash = "sha256:fdab72030b69625665b2eeb9738bdde748ad254e8073085a0ce95382678e8168", size = 426820, upload-time = "2026-03-24T09:56:36.575Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/47/d4/2fdf854bc3b9c7f55219678f812600a20a138af2dd847d99004994eada8f/sentry_sdk-2.53.0-py2.py3-none-any.whl", hash = "sha256:46e1ed8d84355ae54406c924f6b290c3d61f4048625989a723fd622aab838899", size = 437908, upload-time = "2026-02-16T11:11:13.227Z" }, + { url = "https://files.pythonhosted.org/packages/cd/1a/b3a3e9f6520493fed7997af4d2de7965d71549c62f994a8fd15f2ecd519e/sentry_sdk-2.56.0-py2.py3-none-any.whl", hash = "sha256:5afafb744ceb91d22f4cc650c6bd048ac6af5f7412dcc6c59305a2e36f4dbc02", size = 451568, upload-time = "2026-03-24T09:56:34.807Z" }, ] [[package]] name = "setuptools" -version = "82.0.0" +version = "81.0.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/82/f3/748f4d6f65d1756b9ae577f329c951cda23fb900e4de9f70900ced962085/setuptools-82.0.0.tar.gz", hash = "sha256:22e0a2d69474c6ae4feb01951cb69d515ed23728cf96d05513d36e42b62b37cb", size = 1144893, upload-time = "2026-02-08T15:08:40.206Z" } +sdist = { url = "https://files.pythonhosted.org/packages/0d/1c/73e719955c59b8e424d015ab450f51c0af856ae46ea2da83eba51cc88de1/setuptools-81.0.0.tar.gz", hash = "sha256:487b53915f52501f0a79ccfd0c02c165ffe06631443a886740b91af4b7a5845a", size = 1198299, upload-time = "2026-02-06T21:10:39.601Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e1/c6/76dc613121b793286a3f91621d7b75a2b493e0390ddca50f11993eadf192/setuptools-82.0.0-py3-none-any.whl", hash = "sha256:70b18734b607bd1da571d097d236cfcfacaf01de45717d59e6e04b96877532e0", size = 1003468, upload-time = "2026-02-08T15:08:38.723Z" }, + { url = "https://files.pythonhosted.org/packages/e1/e3/c164c88b2e5ce7b24d667b9bd83589cf4f3520d97cad01534cd3c4f55fdb/setuptools-81.0.0-py3-none-any.whl", hash = "sha256:fdd925d5c5d9f62e4b74b30d6dd7828ce236fd6ed998a08d81de62ce5a6310d6", size = 1062021, upload-time = "2026-02-06T21:10:37.175Z" }, ] [[package]] @@ -1619,11 +2143,28 @@ wheels = [ [[package]] name = "smmap" -version = "5.0.2" +version = "5.0.3" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/44/cd/a040c4b3119bbe532e5b0732286f805445375489fceaec1f48306068ee3b/smmap-5.0.2.tar.gz", hash = "sha256:26ea65a03958fa0c8a1c7e8c7a58fdc77221b8910f6be2131affade476898ad5", size = 22329, upload-time = "2025-01-02T07:14:40.909Z" } +sdist = { url = "https://files.pythonhosted.org/packages/1f/ea/49c993d6dfdd7338c9b1000a0f36817ed7ec84577ae2e52f890d1a4ff909/smmap-5.0.3.tar.gz", hash = "sha256:4d9debb8b99007ae47165abc08670bd74cb74b5227dda7f643eccc4e9eb5642c", size = 22506, upload-time = "2026-03-09T03:43:26.1Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/04/be/d09147ad1ec7934636ad912901c5fd7667e1c858e19d355237db0d0cd5e4/smmap-5.0.2-py3-none-any.whl", hash = "sha256:b30115f0def7d7531d22a0fb6502488d879e75b260a9db4d0819cfb25403af5e", size = 24303, upload-time = "2025-01-02T07:14:38.724Z" }, + { url = "https://files.pythonhosted.org/packages/c1/d4/59e74daffcb57a07668852eeeb6035af9f32cbfd7a1d2511f17d2fe6a738/smmap-5.0.3-py3-none-any.whl", hash = "sha256:c106e05d5a61449cf6ba9a1e650227ecfb141590d2a98412103ff35d89fc7b2f", size = 24390, upload-time = "2026-03-09T03:43:24.361Z" }, +] + +[[package]] +name = "stable-baselines3" +version = "2.7.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cloudpickle" }, + { name = "gymnasium" }, + { name = "matplotlib" }, + { name = "numpy" }, + { name = "pandas" }, + { name = "torch" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c9/42/f284c28272422262a99cdf35ecd2e283fded2f75327e6d5e82a9f6d6fe62/stable_baselines3-2.7.1.tar.gz", hash = "sha256:cd90d12d9ee0d9584053f12215c1682b313be4e3a8d8007739319799c3d2c071", size = 220719, upload-time = "2025-12-05T11:22:03.691Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/df/cc/a3038d3833f329dcd03b2dce8b778e4b41044caff88b48429473b8629623/stable_baselines3-2.7.1-py3-none-any.whl", hash = "sha256:b017e76dfe5ca0ce6eabb29e79c42e8c7e125d5862bfcd43ce04ec19732348d0", size = 188039, upload-time = "2025-12-05T11:22:00.819Z" }, ] [[package]] @@ -1640,21 +2181,63 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f1/7b/ce1eafaf1a76852e2ec9b22edecf1daa58175c090266e9f6c64afcd81d91/stack_data-0.6.3-py3-none-any.whl", hash = "sha256:d5558e0c25a4cb0853cddad3d77da9891a08cb85dd9f9f91b9f8cd66e511e695", size = 24521, upload-time = "2023-09-30T13:58:03.53Z" }, ] +[[package]] +name = "sympy" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mpmath" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, +] + +[[package]] +name = "tensorboard" +version = "2.20.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "grpcio" }, + { name = "markdown" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pillow" }, + { name = "protobuf" }, + { name = "setuptools" }, + { name = "tensorboard-data-server" }, + { name = "werkzeug" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/9c/d9/a5db55f88f258ac669a92858b70a714bbbd5acd993820b41ec4a96a4d77f/tensorboard-2.20.0-py3-none-any.whl", hash = "sha256:9dc9f978cb84c0723acf9a345d96c184f0293d18f166bb8d59ee098e6cfaaba6", size = 5525680, upload-time = "2025-07-17T19:20:49.638Z" }, +] + +[[package]] +name = "tensorboard-data-server" +version = "0.7.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7a/13/e503968fefabd4c6b2650af21e110aa8466fe21432cd7c43a84577a89438/tensorboard_data_server-0.7.2-py3-none-any.whl", hash = "sha256:7e0610d205889588983836ec05dc098e80f97b7e7bbff7e994ebb78f578d0ddb", size = 2356, upload-time = "2023-10-23T21:23:32.16Z" }, + { url = "https://files.pythonhosted.org/packages/b7/85/dabeaf902892922777492e1d253bb7e1264cadce3cea932f7ff599e53fea/tensorboard_data_server-0.7.2-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:9fe5d24221b29625dbc7328b0436ca7fc1c23de4acf4d272f1180856e32f9f60", size = 4823598, upload-time = "2023-10-23T21:23:33.714Z" }, + { url = "https://files.pythonhosted.org/packages/73/c6/825dab04195756cf8ff2e12698f22513b3db2f64925bdd41671bfb33aaa5/tensorboard_data_server-0.7.2-py3-none-manylinux_2_31_x86_64.whl", hash = "sha256:ef687163c24185ae9754ed5650eb5bc4d84ff257aabdc33f0cc6f74d8ba54530", size = 6590363, upload-time = "2023-10-23T21:23:35.583Z" }, +] + [[package]] name = "tensorstore" -version = "0.1.81" +version = "0.1.82" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "ml-dtypes" }, { name = "numpy" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/43/f6/e2403fc05b97ba74ad408a98a42c288e6e1b8eacc23780c153b0e5166179/tensorstore-0.1.81.tar.gz", hash = "sha256:687546192ea6f6c8ae28d18f13103336f68017d928b9f5a00325e9b0548d9c25", size = 7120819, upload-time = "2026-02-06T18:56:12.535Z" } +sdist = { url = "https://files.pythonhosted.org/packages/cd/9b/43aedb544937f214dd7c665a7edf1b8b74f2f55d53ebd351c0ce69acf81a/tensorstore-0.1.82.tar.gz", hash = "sha256:ccfceffb7611fc61330f6da24b8b0abd9251d480ac8a5bac5a1729f9ed0c3a9f", size = 7160364, upload-time = "2026-03-13T00:22:16.888Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/85/82/00037db699f74d792efe2696305ddd6932e04306899e3701824a7f7de961/tensorstore-0.1.81-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:7aefa1e3eadca804bce05215184c9cde29205ac2f3b443ca15a4e1846d31af4e", size = 16521245, upload-time = "2026-02-06T18:55:25.559Z" }, - { url = "https://files.pythonhosted.org/packages/86/2e/1deca1b955cb959eec13fd342ffaa2fd84e4770b4e2bcb95a2f541875a52/tensorstore-0.1.81-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:7e001d3edc6758eb5dc80556da9e945c1381f0529102fcc0301358ba6b9b70ed", size = 14543561, upload-time = "2026-02-06T18:55:27.624Z" }, - { url = "https://files.pythonhosted.org/packages/6c/e4/b4343eae773f72a8777f82c5328191a06d8a5195e62105c14b7dcc49823f/tensorstore-0.1.81-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6c27e07f4e91e6dc6a0878e13e2c5931d1716196b67b0df927f2f571de2576e9", size = 19043982, upload-time = "2026-02-06T18:55:30.076Z" }, - { url = "https://files.pythonhosted.org/packages/31/6c/d8c8508a9f4a83dc910d2365c484ba0debf5e531782065e3657fc8fc9b54/tensorstore-0.1.81-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fcb4786c4955e2d88d518b5b5a367427e3ad21d059cba366ad7aebf5fcc2302e", size = 21049171, upload-time = "2026-02-06T18:55:34.383Z" }, - { url = "https://files.pythonhosted.org/packages/44/a9/c1a751e35a0fcff7f795398c4f98b6c8ea0f00fe7d7704f66a1e08d4352f/tensorstore-0.1.81-cp312-cp312-win_amd64.whl", hash = "sha256:b96cbf1ee74d9038762b2d81305ee1589ec89913a440df6cbd514bc5879655d2", size = 13226573, upload-time = "2026-02-06T18:55:36.463Z" }, + { url = "https://files.pythonhosted.org/packages/0d/c3/5ab0b99487b2596bdc0ebd3a569e50415949a63bad90b18e6476de91a7bb/tensorstore-0.1.82-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:f0ac091bd47ea6f051fe11230ad2642c254b46a8fabdd5184b0600556b5529ed", size = 16570668, upload-time = "2026-03-13T00:21:36.386Z" }, + { url = "https://files.pythonhosted.org/packages/aa/95/92b00a4b2e6192528a9c5bac9f53007acf4aa5d54943b9e114bedb72b2da/tensorstore-0.1.82-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8cae7d0c9b2fa0653f90b147daaf9ed04664cab7d297b9772efcfa088da26cab", size = 14904517, upload-time = "2026-03-13T00:21:38.464Z" }, + { url = "https://files.pythonhosted.org/packages/46/7e/c9c8ad65ee4015787e32d31bcf8278fcb27109e809f8334a64285bd73028/tensorstore-0.1.82-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:34c491ea3c6c1904d4618bfe40020bd83aaeb19d52a266ea0f6919eb3fdc64c4", size = 19344428, upload-time = "2026-03-13T00:21:40.575Z" }, + { url = "https://files.pythonhosted.org/packages/f9/8a/590bb60a190d414abd2f83dd5b5148722d0c5d310a73e21b7a60ab98cf00/tensorstore-0.1.82-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d4182300d8ffa172e961e79c6bd89e38ce6bc5cd3abf1a7dacb22c2396ce40b7", size = 20964954, upload-time = "2026-03-13T00:21:42.515Z" }, + { url = "https://files.pythonhosted.org/packages/43/1c/34e6e97426e1718106e9cb74d3045992bdea3ee368f9ea4ea25b809bdba8/tensorstore-0.1.82-cp312-cp312-win_amd64.whl", hash = "sha256:6369809d01edf66cd487cde5c94f57138167c09561f3d906020fd53c72687f92", size = 13393361, upload-time = "2026-03-13T00:21:44.443Z" }, ] [[package]] @@ -1667,22 +2250,47 @@ wheels = [ ] [[package]] -name = "tornado" -version = "6.5.4" +name = "torch" +version = "2.11.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/37/1d/0a336abf618272d53f62ebe274f712e213f5a03c0b2339575430b8362ef2/tornado-6.5.4.tar.gz", hash = "sha256:a22fa9047405d03260b483980635f0b041989d8bcc9a313f8fe18b411d84b1d7", size = 513632, upload-time = "2025-12-15T19:21:03.836Z" } +dependencies = [ + { name = "cuda-bindings", marker = "sys_platform == 'linux'" }, + { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" }, + { name = "setuptools" }, + { name = "sympy" }, + { name = "triton", marker = "sys_platform == 'linux'" }, + { name = "typing-extensions" }, +] wheels = [ - { url = "https://files.pythonhosted.org/packages/ab/a9/e94a9d5224107d7ce3cc1fab8d5dc97f5ea351ccc6322ee4fb661da94e35/tornado-6.5.4-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:d6241c1a16b1c9e4cc28148b1cda97dd1c6cb4fb7068ac1bedc610768dff0ba9", size = 443909, upload-time = "2025-12-15T19:20:48.382Z" }, - { url = "https://files.pythonhosted.org/packages/db/7e/f7b8d8c4453f305a51f80dbb49014257bb7d28ccb4bbb8dd328ea995ecad/tornado-6.5.4-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:2d50f63dda1d2cac3ae1fa23d254e16b5e38153758470e9956cbc3d813d40843", size = 442163, upload-time = "2025-12-15T19:20:49.791Z" }, - { url = "https://files.pythonhosted.org/packages/ba/b5/206f82d51e1bfa940ba366a8d2f83904b15942c45a78dd978b599870ab44/tornado-6.5.4-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d1cf66105dc6acb5af613c054955b8137e34a03698aa53272dbda4afe252be17", size = 445746, upload-time = "2025-12-15T19:20:51.491Z" }, - { url = "https://files.pythonhosted.org/packages/8e/9d/1a3338e0bd30ada6ad4356c13a0a6c35fbc859063fa7eddb309183364ac1/tornado-6.5.4-cp39-abi3-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:50ff0a58b0dc97939d29da29cd624da010e7f804746621c78d14b80238669335", size = 445083, upload-time = "2025-12-15T19:20:52.778Z" }, - { url = "https://files.pythonhosted.org/packages/50/d4/e51d52047e7eb9a582da59f32125d17c0482d065afd5d3bc435ff2120dc5/tornado-6.5.4-cp39-abi3-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e5fb5e04efa54cf0baabdd10061eb4148e0be137166146fff835745f59ab9f7f", size = 445315, upload-time = "2025-12-15T19:20:53.996Z" }, - { url = "https://files.pythonhosted.org/packages/27/07/2273972f69ca63dbc139694a3fc4684edec3ea3f9efabf77ed32483b875c/tornado-6.5.4-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:9c86b1643b33a4cd415f8d0fe53045f913bf07b4a3ef646b735a6a86047dda84", size = 446003, upload-time = "2025-12-15T19:20:56.101Z" }, - { url = "https://files.pythonhosted.org/packages/d1/83/41c52e47502bf7260044413b6770d1a48dda2f0246f95ee1384a3cd9c44a/tornado-6.5.4-cp39-abi3-musllinux_1_2_i686.whl", hash = "sha256:6eb82872335a53dd063a4f10917b3efd28270b56a33db69009606a0312660a6f", size = 445412, upload-time = "2025-12-15T19:20:57.398Z" }, - { url = "https://files.pythonhosted.org/packages/10/c7/bc96917f06cbee182d44735d4ecde9c432e25b84f4c2086143013e7b9e52/tornado-6.5.4-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:6076d5dda368c9328ff41ab5d9dd3608e695e8225d1cd0fd1e006f05da3635a8", size = 445392, upload-time = "2025-12-15T19:20:58.692Z" }, - { url = "https://files.pythonhosted.org/packages/0c/1a/d7592328d037d36f2d2462f4bc1fbb383eec9278bc786c1b111cbbd44cfa/tornado-6.5.4-cp39-abi3-win32.whl", hash = "sha256:1768110f2411d5cd281bac0a090f707223ce77fd110424361092859e089b38d1", size = 446481, upload-time = "2025-12-15T19:21:00.008Z" }, - { url = "https://files.pythonhosted.org/packages/d6/6d/c69be695a0a64fd37a97db12355a035a6d90f79067a3cf936ec2b1dc38cd/tornado-6.5.4-cp39-abi3-win_amd64.whl", hash = "sha256:fa07d31e0cd85c60713f2b995da613588aa03e1303d75705dca6af8babc18ddc", size = 446886, upload-time = "2025-12-15T19:21:01.287Z" }, - { url = "https://files.pythonhosted.org/packages/50/49/8dc3fd90902f70084bd2cd059d576ddb4f8bb44c2c7c0e33a11422acb17e/tornado-6.5.4-cp39-abi3-win_arm64.whl", hash = "sha256:053e6e16701eb6cbe641f308f4c1a9541f91b6261991160391bfc342e8a551a1", size = 445910, upload-time = "2025-12-15T19:21:02.571Z" }, + { url = "https://files.pythonhosted.org/packages/6f/8b/69e3008d78e5cee2b30183340cc425081b78afc5eff3d080daab0adda9aa/torch-2.11.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4b5866312ee6e52ea625cd211dcb97d6a2cdc1131a5f15cc0d87eec948f6dd34", size = 80606338, upload-time = "2026-03-23T18:11:34.781Z" }, + { url = "https://files.pythonhosted.org/packages/13/16/42e5915ebe4868caa6bac83a8ed59db57f12e9a61b7d749d584776ed53d5/torch-2.11.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:f99924682ef0aa6a4ab3b1b76f40dc6e273fca09f367d15a524266db100a723f", size = 419731115, upload-time = "2026-03-23T18:11:06.944Z" }, + { url = "https://files.pythonhosted.org/packages/1a/c9/82638ef24d7877510f83baf821f5619a61b45568ce21c0a87a91576510aa/torch-2.11.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:0f68f4ac6d95d12e896c3b7a912b5871619542ec54d3649cf48cc1edd4dd2756", size = 530712279, upload-time = "2026-03-23T18:10:31.481Z" }, + { url = "https://files.pythonhosted.org/packages/1c/ff/6756f1c7ee302f6d202120e0f4f05b432b839908f9071157302cedfc5232/torch-2.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:fbf39280699d1b869f55eac536deceaa1b60bd6788ba74f399cc67e60a5fab10", size = 114556047, upload-time = "2026-03-23T18:10:55.931Z" }, +] + +[[package]] +name = "tornado" +version = "6.5.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f8/f1/3173dfa4a18db4a9b03e5d55325559dab51ee653763bb8745a75af491286/tornado-6.5.5.tar.gz", hash = "sha256:192b8f3ea91bd7f1f50c06955416ed76c6b72f96779b962f07f911b91e8d30e9", size = 516006, upload-time = "2026-03-10T21:31:02.067Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/59/8c/77f5097695f4dd8255ecbd08b2a1ed8ba8b953d337804dd7080f199e12bf/tornado-6.5.5-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:487dc9cc380e29f58c7ab88f9e27cdeef04b2140862e5076a66fb6bb68bb1bfa", size = 445983, upload-time = "2026-03-10T21:30:44.28Z" }, + { url = "https://files.pythonhosted.org/packages/ab/5e/7625b76cd10f98f1516c36ce0346de62061156352353ef2da44e5c21523c/tornado-6.5.5-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:65a7f1d46d4bb41df1ac99f5fcb685fb25c7e61613742d5108b010975a9a6521", size = 444246, upload-time = "2026-03-10T21:30:46.571Z" }, + { url = "https://files.pythonhosted.org/packages/b2/04/7b5705d5b3c0fab088f434f9c83edac1573830ca49ccf29fb83bf7178eec/tornado-6.5.5-cp39-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:e74c92e8e65086b338fd56333fb9a68b9f6f2fe7ad532645a290a464bcf46be5", size = 447229, upload-time = "2026-03-10T21:30:48.273Z" }, + { url = "https://files.pythonhosted.org/packages/34/01/74e034a30ef59afb4097ef8659515e96a39d910b712a89af76f5e4e1f93c/tornado-6.5.5-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:435319e9e340276428bbdb4e7fa732c2d399386d1de5686cb331ec8eee754f07", size = 448192, upload-time = "2026-03-10T21:30:51.22Z" }, + { url = "https://files.pythonhosted.org/packages/be/00/fe9e02c5a96429fce1a1d15a517f5d8444f9c412e0bb9eadfbe3b0fc55bf/tornado-6.5.5-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:3f54aa540bdbfee7b9eb268ead60e7d199de5021facd276819c193c0fb28ea4e", size = 448039, upload-time = "2026-03-10T21:30:53.52Z" }, + { url = "https://files.pythonhosted.org/packages/82/9e/656ee4cec0398b1d18d0f1eb6372c41c6b889722641d84948351ae19556d/tornado-6.5.5-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:36abed1754faeb80fbd6e64db2758091e1320f6bba74a4cf8c09cd18ccce8aca", size = 447445, upload-time = "2026-03-10T21:30:55.541Z" }, + { url = "https://files.pythonhosted.org/packages/5a/76/4921c00511f88af86a33de770d64141170f1cfd9c00311aea689949e274e/tornado-6.5.5-cp39-abi3-win32.whl", hash = "sha256:dd3eafaaeec1c7f2f8fdcd5f964e8907ad788fe8a5a32c4426fbbdda621223b7", size = 448582, upload-time = "2026-03-10T21:30:57.142Z" }, + { url = "https://files.pythonhosted.org/packages/2c/23/f6c6112a04d28eed765e374435fb1a9198f73e1ec4b4024184f21faeb1ad/tornado-6.5.5-cp39-abi3-win_amd64.whl", hash = "sha256:6443a794ba961a9f619b1ae926a2e900ac20c34483eea67be4ed8f1e58d3ef7b", size = 448990, upload-time = "2026-03-10T21:30:58.857Z" }, + { url = "https://files.pythonhosted.org/packages/b7/c8/876602cbc96469911f0939f703453c1157b0c826ecb05bdd32e023397d4e/tornado-6.5.5-cp39-abi3-win_arm64.whl", hash = "sha256:2c9a876e094109333f888539ddb2de4361743e5d21eece20688e3e351e4990a6", size = 448016, upload-time = "2026-03-10T21:31:00.43Z" }, ] [[package]] @@ -1732,14 +2340,23 @@ wheels = [ [[package]] name = "trimesh" -version = "4.11.2" +version = "4.11.4" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/1f/41/de14e2fa9b2d99214c60402fc57d2efb201f2925b16d6bee289565901d83/trimesh-4.11.2.tar.gz", hash = "sha256:30fbde5b8dd7c157e7ff4d54286cb35291844fd3f4d0364e8b2727f1b308fb06", size = 835044, upload-time = "2026-02-10T16:00:27.58Z" } +sdist = { url = "https://files.pythonhosted.org/packages/1c/6c/57a77091f42c4fe3246810c8878b1f08c65944432bb856e1b797e960c822/trimesh-4.11.4.tar.gz", hash = "sha256:9c3bf253f8b21978e905c2f2fa361621415a6dfaac6b7fdaa54ef3f7f66b8c79", size = 836069, upload-time = "2026-03-18T22:59:11.357Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3c/b9/da09903ea53b677a58ba770112de6fe8b2acb8b4cd9bffae4ff6cfe7c072/trimesh-4.11.2-py3-none-any.whl", hash = "sha256:25e3ab2620f9eca5c9376168c67aabdd32205dad1c4eea09cd45cd4a3edf775a", size = 740328, upload-time = "2026-02-10T16:00:25.246Z" }, + { url = "https://files.pythonhosted.org/packages/74/3a/0b9fb22a6c34cff36d70d1eb83bf61540aa2d7ced0f5ee023eb2123c3aa2/trimesh-4.11.4-py3-none-any.whl", hash = "sha256:7606a3be929ced36a3bbda8044d675510c46f83fe675fd9a354b5cf13f7db7ae", size = 740767, upload-time = "2026-03-18T22:59:09.45Z" }, +] + +[[package]] +name = "triton" +version = "3.6.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/17/5d/08201db32823bdf77a0e2b9039540080b2e5c23a20706ddba942924ebcd6/triton-3.6.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:374f52c11a711fd062b4bfbb201fd9ac0a5febd28a96fb41b4a0f51dde3157f4", size = 176128243, upload-time = "2026-01-20T16:16:07.857Z" }, + { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, ] [[package]] @@ -1751,6 +2368,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, ] +[[package]] +name = "typing-inspect" +version = "0.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mypy-extensions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/dc/74/1789779d91f1961fa9438e9a8710cdae6bd138c80d7303996933d117264a/typing_inspect-0.9.0.tar.gz", hash = "sha256:b23fc42ff6f6ef6954e4852c1fb512cdd18dbea03134f91f856a95ccc9461f78", size = 13825, upload-time = "2023-05-24T20:25:47.612Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/65/f3/107a22063bf27bdccf2024833d3445f4eea42b2e598abfbd46f6a63b6cb0/typing_inspect-0.9.0-py3-none-any.whl", hash = "sha256:9ee6fc59062311ef8547596ab6b955e1b8aa46242d854bfc78f4f6b0eff35f9f", size = 8827, upload-time = "2023-05-24T20:25:45.287Z" }, +] + [[package]] name = "typing-inspection" version = "0.4.2" @@ -1763,6 +2393,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, ] +[[package]] +name = "tzdata" +version = "2025.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5e/a7/c202b344c5ca7daf398f3b8a477eeb205cf3b6f32e7ec3a6bac0629ca975/tzdata-2025.3.tar.gz", hash = "sha256:de39c2ca5dc7b0344f2eba86f49d614019d29f060fc4ebc8a417896a620b56a7", size = 196772, upload-time = "2025-12-13T17:45:35.667Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c7/b0/003792df09decd6849a5e39c28b513c06e84436a54440380862b5aeff25d/tzdata-2025.3-py2.py3-none-any.whl", hash = "sha256:06a47e5700f3081aab02b2e513160914ff0694bce9947d6b76ebd6bf57cfc5d1", size = 348521, upload-time = "2025-12-13T17:45:33.889Z" }, +] + [[package]] name = "urllib3" version = "2.6.3" @@ -1786,6 +2425,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/99/39/6b3f7d234ba3964c428a6e40006340f53ba37993f46ed6e111c6e9141d18/uvloop-0.22.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:512fec6815e2dd45161054592441ef76c830eddaad55c8aa30952e6fe1ed07c0", size = 4296343, upload-time = "2025-10-16T22:16:35.149Z" }, ] +[[package]] +name = "virtualenv" +version = "21.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "distlib" }, + { name = "filelock" }, + { name = "platformdirs" }, + { name = "python-discovery" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/aa/92/58199fe10049f9703c2666e809c4f686c54ef0a68b0f6afccf518c0b1eb9/virtualenv-21.2.0.tar.gz", hash = "sha256:1720dc3a62ef5b443092e3f499228599045d7fea4c79199770499df8becf9098", size = 5840618, upload-time = "2026-03-09T17:24:38.013Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c6/59/7d02447a55b2e55755011a647479041bc92a82e143f96a8195cb33bd0a1c/virtualenv-21.2.0-py3-none-any.whl", hash = "sha256:1bd755b504931164a5a496d217c014d098426cddc79363ad66ac78125f9d908f", size = 5825084, upload-time = "2026-03-09T17:24:35.378Z" }, +] + +[[package]] +name = "wadler-lindig" +version = "0.1.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1e/67/cbae4bf7683a64755c2c1778c418fea96d00e34395bb91743f08bd951571/wadler_lindig-0.1.7.tar.gz", hash = "sha256:81d14d3fe77d441acf3ebd7f4aefac20c74128bf460e84b512806dccf7b2cd55", size = 15842, upload-time = "2025-06-18T07:00:42.843Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/96/04e7b441807b26b794da5b11e59ed7f83b2cf8af202bd7eba8ad2fa6046e/wadler_lindig-0.1.7-py3-none-any.whl", hash = "sha256:e3ec83835570fd0a9509f969162aeb9c65618f998b1f42918cfc8d45122fe953", size = 20516, upload-time = "2025-06-18T07:00:41.684Z" }, +] + [[package]] name = "wandb" version = "0.24.2" @@ -1815,6 +2478,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3a/9a/f3919d7ee7ba99dabf0aac7e299c6c328f5eae94f9f6b28c76005f882d5d/wandb-0.24.2-py3-none-win_arm64.whl", hash = "sha256:b42614b99f8b9af69f88c15a84283a973c8cd5750e9c4752aa3ce21f13dbac9a", size = 20268261, upload-time = "2026-02-05T00:12:14.353Z" }, ] +[[package]] +name = "warp-lang" +version = "1.12.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/7d/15/fadf3e3ba5c1c907530c20c98402aaef792da74bbbe382c848cef6e5affe/warp_lang-1.12.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:c78c3701d5cad86c30ef5017410d294ec46a396bb0d502ee1c98743494f3a62f", size = 24168341, upload-time = "2026-03-06T19:42:16.333Z" }, + { url = "https://files.pythonhosted.org/packages/98/13/deab9dbae5c6aa753ac8ea1d3b1f85d20c5bab7bdebd8916ce242fbe1f0b/warp_lang-1.12.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:a1436f60a1881cd94f787e751a83fc0987626be2d3e2b4e74c64a6947c6d1266", size = 136485344, upload-time = "2026-03-06T19:43:02.427Z" }, + { url = "https://files.pythonhosted.org/packages/45/ce/9f5c57cac849edaba2f3335cb649b7019b09195b3af02221258482254559/warp_lang-1.12.0-py3-none-manylinux_2_34_aarch64.whl", hash = "sha256:a2d6decba693aba5b828573c4414fd6a3f4c4a934db9c322736ef2b3fa99fe76", size = 137735580, upload-time = "2026-03-06T19:44:22.279Z" }, + { url = "https://files.pythonhosted.org/packages/7e/3f/1ddc888fe769447ae33915a9567a9dd7467e1fc7fc8010d39e01b339667f/warp_lang-1.12.0-py3-none-win_amd64.whl", hash = "sha256:697248edd2f1e2952f50e3db33b214af76173641a8894aacc467bed6dc247f8a", size = 119793582, upload-time = "2026-03-06T19:45:37.288Z" }, +] + [[package]] name = "wcwidth" version = "0.6.0" @@ -1825,21 +2502,35 @@ wheels = [ ] [[package]] -name = "wrapt" -version = "2.1.1" +name = "werkzeug" +version = "3.1.7" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/f7/37/ae31f40bec90de2f88d9597d0b5281e23ffe85b893a47ca5d9c05c63a4f6/wrapt-2.1.1.tar.gz", hash = "sha256:5fdcb09bf6db023d88f312bd0767594b414655d58090fc1c46b3414415f67fac", size = 81329, upload-time = "2026-02-03T02:12:13.786Z" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b5/43/76ded108b296a49f52de6bac5192ca1c4be84e886f9b5c9ba8427d9694fd/werkzeug-3.1.7.tar.gz", hash = "sha256:fb8c01fe6ab13b9b7cdb46892b99b1d66754e1d7ab8e542e865ec13f526b5351", size = 875700, upload-time = "2026-03-24T01:08:07.687Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/df/cb/4d5255d19bbd12be7f8ee2c1fb4269dddec9cef777ef17174d357468efaa/wrapt-2.1.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:ab8e3793b239db021a18782a5823fcdea63b9fe75d0e340957f5828ef55fcc02", size = 61143, upload-time = "2026-02-03T02:11:46.313Z" }, - { url = "https://files.pythonhosted.org/packages/6f/07/7ed02daa35542023464e3c8b7cb937fa61f6c61c0361ecf8f5fecf8ad8da/wrapt-2.1.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:7c0300007836373d1c2df105b40777986accb738053a92fe09b615a7a4547e9f", size = 61740, upload-time = "2026-02-03T02:12:51.966Z" }, - { url = "https://files.pythonhosted.org/packages/c4/60/a237a4e4a36f6d966061ccc9b017627d448161b19e0a3ab80a7c7c97f859/wrapt-2.1.1-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:2b27c070fd1132ab23957bcd4ee3ba707a91e653a9268dc1afbd39b77b2799f7", size = 121327, upload-time = "2026-02-03T02:11:06.796Z" }, - { url = "https://files.pythonhosted.org/packages/ae/fe/9139058a3daa8818fc67e6460a2340e8bbcf3aef8b15d0301338bbe181ca/wrapt-2.1.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8b0e36d845e8b6f50949b6b65fc6cd279f47a1944582ed4ec8258cd136d89a64", size = 122903, upload-time = "2026-02-03T02:12:48.657Z" }, - { url = "https://files.pythonhosted.org/packages/91/10/b8479202b4164649675846a531763531f0a6608339558b5a0a718fc49a8d/wrapt-2.1.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:4aeea04a9889370fcfb1ef828c4cc583f36a875061505cd6cd9ba24d8b43cc36", size = 121333, upload-time = "2026-02-03T02:11:32.148Z" }, - { url = "https://files.pythonhosted.org/packages/5f/75/75fc793b791d79444aca2c03ccde64e8b99eda321b003f267d570b7b0985/wrapt-2.1.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:d88b46bb0dce9f74b6817bc1758ff2125e1ca9e1377d62ea35b6896142ab6825", size = 120458, upload-time = "2026-02-03T02:11:16.039Z" }, - { url = "https://files.pythonhosted.org/packages/d7/8f/c3f30d511082ca6d947c405f9d8f6c8eaf83cfde527c439ec2c9a30eb5ea/wrapt-2.1.1-cp312-cp312-win32.whl", hash = "sha256:63decff76ca685b5c557082dfbea865f3f5f6d45766a89bff8dc61d336348833", size = 58086, upload-time = "2026-02-03T02:12:35.041Z" }, - { url = "https://files.pythonhosted.org/packages/0a/c8/37625b643eea2849f10c3b90f69c7462faa4134448d4443234adaf122ae5/wrapt-2.1.1-cp312-cp312-win_amd64.whl", hash = "sha256:b828235d26c1e35aca4107039802ae4b1411be0fe0367dd5b7e4d90e562fcbcd", size = 60328, upload-time = "2026-02-03T02:12:45.808Z" }, - { url = "https://files.pythonhosted.org/packages/ce/79/56242f07572d5682ba8065a9d4d9c2218313f576e3c3471873c2a5355ffd/wrapt-2.1.1-cp312-cp312-win_arm64.whl", hash = "sha256:75128507413a9f1bcbe2db88fd18fbdbf80f264b82fa33a6996cdeaf01c52352", size = 58722, upload-time = "2026-02-03T02:12:27.949Z" }, - { url = "https://files.pythonhosted.org/packages/c4/da/5a086bf4c22a41995312db104ec2ffeee2cf6accca9faaee5315c790377d/wrapt-2.1.1-py3-none-any.whl", hash = "sha256:3b0f4629eb954394a3d7c7a1c8cca25f0b07cefe6aa8545e862e9778152de5b7", size = 43886, upload-time = "2026-02-03T02:11:45.048Z" }, + { url = "https://files.pythonhosted.org/packages/7f/b2/0bba9bbb4596d2d2f285a16c2ab04118f6b957d8441566e1abb892e6a6b2/werkzeug-3.1.7-py3-none-any.whl", hash = "sha256:4b314d81163a3e1a169b6a0be2a000a0e204e8873c5de6586f453c55688d422f", size = 226295, upload-time = "2026-03-24T01:08:06.133Z" }, +] + +[[package]] +name = "wrapt" +version = "2.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2e/64/925f213fdcbb9baeb1530449ac71a4d57fc361c053d06bf78d0c5c7cd80c/wrapt-2.1.2.tar.gz", hash = "sha256:3996a67eecc2c68fd47b4e3c564405a5777367adfd9b8abb58387b63ee83b21e", size = 81678, upload-time = "2026-03-06T02:53:25.134Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4c/b6/1db817582c49c7fcbb7df6809d0f515af29d7c2fbf57eb44c36e98fb1492/wrapt-2.1.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:ff2aad9c4cda28a8f0653fc2d487596458c2a3f475e56ba02909e950a9efa6a9", size = 61255, upload-time = "2026-03-06T02:52:45.663Z" }, + { url = "https://files.pythonhosted.org/packages/a2/16/9b02a6b99c09227c93cd4b73acc3678114154ec38da53043c0ddc1fba0dc/wrapt-2.1.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6433ea84e1cfacf32021d2a4ee909554ade7fd392caa6f7c13f1f4bf7b8e8748", size = 61848, upload-time = "2026-03-06T02:53:48.728Z" }, + { url = "https://files.pythonhosted.org/packages/af/aa/ead46a88f9ec3a432a4832dfedb84092fc35af2d0ba40cd04aea3889f247/wrapt-2.1.2-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:c20b757c268d30d6215916a5fa8461048d023865d888e437fab451139cad6c8e", size = 121433, upload-time = "2026-03-06T02:54:40.328Z" }, + { url = "https://files.pythonhosted.org/packages/3a/9f/742c7c7cdf58b59085a1ee4b6c37b013f66ac33673a7ef4aaed5e992bc33/wrapt-2.1.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:79847b83eb38e70d93dc392c7c5b587efe65b3e7afcc167aa8abd5d60e8761c8", size = 123013, upload-time = "2026-03-06T02:53:26.58Z" }, + { url = "https://files.pythonhosted.org/packages/e8/44/2c3dd45d53236b7ed7c646fcf212251dc19e48e599debd3926b52310fafb/wrapt-2.1.2-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f8fba1bae256186a83d1875b2b1f4e2d1242e8fac0f58ec0d7e41b26967b965c", size = 117326, upload-time = "2026-03-06T02:53:11.547Z" }, + { url = "https://files.pythonhosted.org/packages/74/e2/b17d66abc26bd96f89dec0ecd0ef03da4a1286e6ff793839ec431b9fae57/wrapt-2.1.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e3d3b35eedcf5f7d022291ecd7533321c4775f7b9cd0050a31a68499ba45757c", size = 121444, upload-time = "2026-03-06T02:54:09.5Z" }, + { url = "https://files.pythonhosted.org/packages/3c/62/e2977843fdf9f03daf1586a0ff49060b1b2fc7ff85a7ea82b6217c1ae36e/wrapt-2.1.2-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:6f2c5390460de57fa9582bc8a1b7a6c86e1a41dfad74c5225fc07044c15cc8d1", size = 116237, upload-time = "2026-03-06T02:54:03.884Z" }, + { url = "https://files.pythonhosted.org/packages/88/dd/27fc67914e68d740bce512f11734aec08696e6b17641fef8867c00c949fc/wrapt-2.1.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:7dfa9f2cf65d027b951d05c662cc99ee3bd01f6e4691ed39848a7a5fffc902b2", size = 120563, upload-time = "2026-03-06T02:53:20.412Z" }, + { url = "https://files.pythonhosted.org/packages/ec/9f/b750b3692ed2ef4705cb305bd68858e73010492b80e43d2a4faa5573cbe7/wrapt-2.1.2-cp312-cp312-win32.whl", hash = "sha256:eba8155747eb2cae4a0b913d9ebd12a1db4d860fc4c829d7578c7b989bd3f2f0", size = 58198, upload-time = "2026-03-06T02:53:37.732Z" }, + { url = "https://files.pythonhosted.org/packages/8e/b2/feecfe29f28483d888d76a48f03c4c4d8afea944dbee2b0cd3380f9df032/wrapt-2.1.2-cp312-cp312-win_amd64.whl", hash = "sha256:1c51c738d7d9faa0b3601708e7e2eda9bf779e1b601dce6c77411f2a1b324a63", size = 60441, upload-time = "2026-03-06T02:52:47.138Z" }, + { url = "https://files.pythonhosted.org/packages/44/e1/e328f605d6e208547ea9fd120804fcdec68536ac748987a68c47c606eea8/wrapt-2.1.2-cp312-cp312-win_arm64.whl", hash = "sha256:c8e46ae8e4032792eb2f677dbd0d557170a8e5524d22acc55199f43efedd39bf", size = 58836, upload-time = "2026-03-06T02:53:22.053Z" }, + { url = "https://files.pythonhosted.org/packages/1a/c7/8528ac2dfa2c1e6708f647df7ae144ead13f0a31146f43c7264b4942bf12/wrapt-2.1.2-py3-none-any.whl", hash = "sha256:b8fd6fa2b2c4e7621808f8c62e8317f4aae56e59721ad933bac5239d913cf0e8", size = 43993, upload-time = "2026-03-06T02:53:12.905Z" }, ] [[package]]