diff --git a/.github/workflows/deps-sync.yml b/.github/workflows/deps-sync.yml new file mode 100644 index 0000000..146439f --- /dev/null +++ b/.github/workflows/deps-sync.yml @@ -0,0 +1,27 @@ +name: Dependencies sync +on: + push: + paths: + - requirements.txt + - pyproject.toml + pull_request: + paths: + - requirements.txt + - pyproject.toml +permissions: + contents: read +jobs: + check-requirements-sync: + name: Check requirements.txt vs pyproject.toml + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v5 + with: + persist-credentials: false + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: '3.12' + - name: Check dependency sync + run: python scripts/check_requirements_sync.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..52815c1 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,38 @@ +name: CI +on: + push: + branches: + - "**" + pull_request: + branches: + - "**" + types: [opened, synchronize, edited, reopened, labeled, ready_for_review] +jobs: + tests: + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [ubuntu-latest, windows-latest] + python-version: ["3.12"] + steps: + - uses: actions/checkout@v5 + - name: Install uv + uses: astral-sh/setup-uv@v7 + with: + enable-cache: true + python-version: ${{ matrix.python-version }} + - name: Install the project + run: uv sync --group dev + - name: Install headless MuJoCo rendering libraries + if: runner.os == 'Linux' + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends libosmesa6 libgl1 libglib2.0-0 + - name: Run tests + run: uv run pytest -v + env: + # Ubuntu runners have no GPU. OSMesa renders in software. EGL fails + # here because libEGL is not installed and PyOpenGL's EGL platform + # stays None. Windows uses MuJoCo's default GLFW backend. + MUJOCO_GL: ${{ runner.os == 'Linux' && 'osmesa' || '' }} + PYOPENGL_PLATFORM: ${{ runner.os == 'Linux' && 'osmesa' || '' }} diff --git a/.gitignore b/.gitignore index 8610959..24f76bd 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,32 @@ -*__pycache__/ +# Python +__pycache__/ +*.py[cod] +*.egg-info/ +build/ +dist/ +.eggs/ + +# Virtual environments +.venv/ +venv/ +env/ + +# Tooling caches +.pytest_cache/ +.ruff_cache/ +.mypy_cache/ +.ipynb_checkpoints/ + +# Experiment outputs +videos/ +logdir/ runs/ -logdir/ \ No newline at end of file +tb/ +tmp/ +*.mp4 + +# OS / editor +.DS_Store +Thumbs.db +.idea/ +.vscode/ diff --git a/.python-version b/.python-version new file mode 100644 index 0000000..e4fba21 --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.12 diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..671e956 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,3 @@ +include requirements.txt +include conda-environment.yaml +include bpo/configs.yaml diff --git a/README.md b/README.md index d2d5b24..89d58fb 100644 --- a/README.md +++ b/README.md @@ -4,117 +4,188 @@ The repository provide the official implementation of **Behaviour Policy Optimiz > **Behaviour Policy Optimization: Provably Lower Variance Return Estimates for Off-Policy Reinforcement Learning** > Alexander W. Goodall, Edwin Hamel-De le Court, Francesco Belardinelli -> arXiv: [https://arxiv.org/abs/2511.10843](https://arxiv.org/abs/2511.10843) +> arXiv: [https://arxiv.org/abs/2511.10843](https://arxiv.org/abs/2511.10843) *Disclaimer:* This implementation is simplified from the full version used to collect the results provided in the paper. +## Software stack + +Training runs the actor-critic on CPU or CUDA GPUs with [JAX](https://docs.jax.dev/), uses [TensorFlow](https://www.tensorflow.org/) and [TensorFlow Probability](https://www.tensorflow.org/probability) for action distributions and TensorBoard, and steps through the environments with [Gymnasium](https://gymnasium.farama.org/). + ## Quick Start ### Installation -#### Prequisites +#### Prerequisites -Python 3.10 is required. +Python 3.12 is required. + +#### Installation with conda -#### Installation with conda - Install conda, e.g., via [anaconda](https://anaconda.org/channels/anaconda/packages/conda/overview). - Clone the repo: + ```bash git clone https://github.com/sacktock/BPO.git cd BPO ``` + - Create a conda virtual environment: + ```bash conda env create --name bpo --file conda-environment.yaml conda activate bpo ``` + - Install the project and its dependencies: + ```bash -pip install -r requirements.txt +pip install -e . ``` -### Check the installation is working +> `pip install -e .` installs the dependencies declared in `pyproject.toml`. `requirements.txt` lists the same base dependencies and is kept in lockstep with that file. -``` -python run.py --configs cartpole_ppo +#### Installation with uv + +From the repository root, `uv sync` installs the base dependency set (CPU JAX): + +```bash +uv sync ``` -### Enabling GPU Acceleration with JAX (Optional) +### Enabling GPU acceleration with JAX (optional) -Our implementation relies on [JAX](https://docs.jax.dev/) for GPU acceleration, which can be enabled for Linux or WSL sub systems. +BPO uses [JAX](https://docs.jax.dev/) for the actor-critic. CPU is enough to smoke test the environments and run a short training job. + +**macOS and Windows:** we recommend JAX with CPU. No further action is required if you correctly followed the earlier steps. + +#### pip + +**Linux x86_64/aarch64** and **Windows WSL2 x86_64**: We strongly recommend using [Ubuntu 22.04](https://apps.microsoft.com/detail/9pn20msr04dw?hl=en-GB&gl=BE) or similar. jax and jaxlib `0.4.30` should already be installed. You need to reinstall [JAX](https://docs.jax.dev/) for your CUDA driver. Do not use the `-U` option here as it may break earlier dependencies. -- **Linux x86_64/aarch64**: jax and jaxlib `0.4.30` should already be installed via the `requirements.txt`. You need to reinstall JAX based on your cuda driver compatibility. Do not use the ```-U``` option here! ```bash pip install "jax[cuda12]" ``` -For `13+` cuda versions you may need to upgrade the jax and jaxlib installation. -- **Windows**: GPU acceletartion is also supported (experimentally) on Windows WSL x86_64. We strongly recommend using [Ubuntu 22.04](https://apps.microsoft.com/detail/9pn20msr04dw?hl=en-GB&gl=BE) or similar. Follow the **Linux x86_64/aarch64** instructions above. +CUDA 13 does not ship with this stack. JAX 0.4.30 only publishes a CUDA 12 extra, so CUDA 13 needs a JAX and jaxlib upgrade, which is a later change. + +#### uv + +```bash +uv sync --group cuda12 +``` + +### Follow-up: native JAX distributions + +TensorFlow Probability pins this repository to TensorFlow 2.17, NumPy 1.26, and JAX 0.4.30. Moving to a newer JAX 0.7 stack, with a native JAX distributions implementation in place of TensorFlow Probability, is a follow-up. It does not ship today, so results stay reproducible on the stack used for the paper. + +### Extras + +MuJoCo (used by Ant, HalfCheetah, Hopper, and Walker2d) is installed through the `gymnasium[mujoco]` extra. On a headless Linux machine (CI job), also install OSMesa and select it, or those environments fail while importing the renderer: + +```bash +sudo apt-get install -y libosmesa6 libgl1 libglib2.0-0 +export MUJOCO_GL=osmesa PYOPENGL_PLATFORM=osmesa +``` + +## How to run BPO + +Training is configured from `bpo/configs.yaml`. Named blocks are merged left to right onto `defaults`. Later configs override earlier ones. -- **MAC**: we recommend JAX with CPU. No further action is required if you correctly followed the earlier steps. +After `pip install -e .` (conda or pip), the `bpo` console script is on `PATH`. With uv, pass the same arguments to `uv run bpo`. -## Reproducing our Results +```bash +bpo --configs cartpole_ppo +``` -All experiments are launched from the command line via `run.py`. Under the hood, `run.py` loads **named configuration blocks** from `configs.yaml` and merges them **left -> right** in the order you pass them to `--configs`. Later configs override earlier ones. +```bash +uv run bpo --configs cartpole_ppo +``` + +Any key in the loaded config can be overridden from the command line (for example `--env.env_id ant`, `--run.seed 0`). ### Running experiments The basic pattern is: ```bash -python run.py --configs [ ...] [flag overrides...] +# conda or pip +bpo --configs [ ...] [flag overrides...] + +# uv +uv run bpo --configs [ ...] [flag overrides...] ``` -* `--configs` can take **one or more** names, each corresponding to a top-level key in `configs.yaml` (e.g., `cartpole_ppo`, `mujoco_ppo_gsde`, `ppo_bpo_zero`, ...). -* After config merging, any remaining CLI arguments are parsed as **overrides** (e.g., `--env.env_id ant`, `--run.seed 0`, etc.). +* `--configs` can take **one or more** names, each corresponding to a top-level key in `bpo/configs.yaml` (e.g., `cartpole_ppo`, `mujoco_ppo_gsde`, `ppo_bpo_zero`, ...). +* After config merging, any remaining CLI arguments are parsed as **overrides** (e.g., `--env.env_id ant`, `--run.seed 0`, etc.). A minimal PPO run (CartPole preset): ```bash -python run.py --configs cartpole_ppo +# conda or pip +bpo --configs cartpole_ppo + +# uv +uv run bpo --configs cartpole_ppo ``` -This uses the `cartpole_ppo` block (timesteps, env, PPO hyperparameters, etc.). +This uses the `cartpole_ppo` block (timesteps, env, PPO hyperparameters, etc.). A MuJoCo PPO run (Ant baseline preset): ```bash -python run.py --configs mujoco_ppo --env.env_id ant --run.seed 0 --run.logdir runs/mujoco/ant/ppo_seed_0 +# conda or pip +bpo --configs mujoco_ppo --env.env_id ant --run.seed 0 --run.logdir runs/mujoco/ant/ppo_seed_0 + +# uv +uv run bpo --configs mujoco_ppo --env.env_id ant --run.seed 0 --run.logdir runs/mujoco/ant/ppo_seed_0 ``` -(`mujoco_ppo` sets MuJoCo-style PPO defaults; `--env.env_id` can be changed among `ant`, `half_cheetah`, `hopper`, `walker_2d`.) +(`mujoco_ppo` sets MuJoCo-style PPO defaults; `--env.env_id` can be changed among `ant`, `half_cheetah`, `hopper`, `walker_2d`.) ### Enabling BPO **Option A -> flip the flag directly** ```bash -python run.py --configs mujoco_ppo_gsde --bpo True --env.env_id ant --run.seed 0 +# conda or pip +bpo --configs mujoco_ppo_gsde --bpo True --env.env_id ant --run.seed 0 + +# uv +uv run bpo --configs mujoco_ppo_gsde --bpo True --env.env_id ant --run.seed 0 ``` **Option B -> include a BPO config block** -For example, `ppo_bpo_zero` is a small "add-on" config that sets `bpo: True` and applies BPO-specific settings (e.g., `symlog_targets`, `polyak_tau`, and a zero-norm-final Q-head). +For example, `ppo_bpo_zero` is a small "add-on" config that sets `bpo: True` and applies BPO-specific settings (e.g., `symlog_targets`, `polyak_tau`, and a zero-norm-final Q-head). ```bash -python run.py --configs mujoco_ppo_gsde ppo_bpo_zero --env.env_id ant --run.seed 0 +# conda or pip +bpo --configs mujoco_ppo_gsde ppo_bpo_zero --env.env_id ant --run.seed 0 + +# uv +uv run bpo --configs mujoco_ppo_gsde ppo_bpo_zero --env.env_id ant --run.seed 0 ``` ### Switching BPO hyperparameter settings -The repo includes several ready-made BPO variants that mainly differ in **importance-weight clipping** (`clip_rho`, `clip_c`) and whether **trajectory clipping** is enabled (`clip_traj`). +The repo includes several ready-made BPO variants that mainly differ in **importance-weight clipping** (`clip_rho`, `clip_c`) and whether **trajectory clipping** is enabled (`clip_traj`). Common presets: -* `ppo_bpo_zero` -> `clip_rho=1.5`, `clip_c=1.5`, plus "zero final norm" Q-head -* `ppo_bpo_zero1.0` -> `clip_rho=1.0`, `clip_c=1.0` -* `ppo_bpo_zero1.0_1.5` -> `clip_rho=1.5`, `clip_c=1.0` -* `ppo_bpo_zero1.0_1.4` -> `clip_rho=1.4`, `clip_c=1.0` -* `ppo_bpo_zero1.0_traj` -> `clip_rho=1.0`, `clip_c=1.0`, `clip_traj=True` -* `ppo_bpo_zero1.5_traj` -> `clip_rho=1.5`, `clip_c=1.5`, `clip_traj=True` +* `ppo_bpo_zero` -> `clip_rho=1.5`, `clip_c=1.5`, plus "zero final norm" Q-head +* `ppo_bpo_zero1.0` -> `clip_rho=1.0`, `clip_c=1.0` +* `ppo_bpo_zero1.0_1.5` -> `clip_rho=1.5`, `clip_c=1.0` +* `ppo_bpo_zero1.0_1.4` -> `clip_rho=1.4`, `clip_c=1.0` +* `ppo_bpo_zero1.0_traj` -> `clip_rho=1.0`, `clip_c=1.0`, `clip_traj=True` +* `ppo_bpo_zero1.5_traj` -> `clip_rho=1.5`, `clip_c=1.5`, `clip_traj=True` So, to run the "rho/c = 1.0/1.0" setting you can do: ```bash -python run.py --configs mujoco_ppo_gsde ppo_bpo_zero1.0 --env.env_id ant --run.seed 0 +# conda or pip +bpo --configs mujoco_ppo_gsde ppo_bpo_zero1.0 --env.env_id ant --run.seed 0 + +# uv +uv run bpo --configs mujoco_ppo_gsde ppo_bpo_zero1.0 --env.env_id ant --run.seed 0 ``` ### Logging and Miscellaneous @@ -122,7 +193,16 @@ python run.py --configs mujoco_ppo_gsde ppo_bpo_zero1.0 --env.env_id ant --run.s **TensorBoard** logging is supported by our implementation and is enabled in the following example: ```bash -python run.py \ +# conda or pip +bpo \ + --configs mujoco_ppo_gsde ppo_bpo_zero \ + --env.env_id ant \ + --tensorboard True \ + --logdir "runs/mujoco/ant/ppo_gsde_fqe_zero_seed_0" \ + --seed 0 + +# uv +uv run bpo \ --configs mujoco_ppo_gsde ppo_bpo_zero \ --env.env_id ant \ --tensorboard True \ @@ -147,17 +227,46 @@ The verbosity level can be changed via a command-line override: Similarly, you can also override any nested field directly from the command line (after configs are merged). For example, to keep `ppo_bpo_zero` but change the clipping thresholds: ```bash -python run.py --configs mujoco_ppo_gsde ppo_bpo_zero \ +# conda or pip +bpo --configs mujoco_ppo_gsde ppo_bpo_zero \ + --env.env_id ant \ + --ppo.clip_rho 1.0 --ppo.clip_c 1.0 \ + --run.seed 0 + +# uv +uv run bpo --configs mujoco_ppo_gsde ppo_bpo_zero \ --env.env_id ant \ --ppo.clip_rho 1.0 --ppo.clip_c 1.0 \ --run.seed 0 ``` +## Tests +With uv: -## License +```bash +uv sync --group dev +uv run pytest -v +``` -Our implementation of BPO is released under the MIT License. +With conda or pip, install the project, then pytest, and run from the repository root: + +```bash +pip install -e . +pip install "pytest>=9" +pytest -v +``` +The test suite is small and checks that the scientific stack imports, that each environment can be constructed and stepped, and that a tiny PPO or BPO training job can start. +## Dependency sync +`requirements.txt` must match `[project.dependencies]` in `pyproject.toml`. CI runs `scripts/check_requirements_sync.py` when either file changes. Locally: + +```bash +python scripts/check_requirements_sync.py +``` + +## License + +Our implementation of BPO is released under the MIT License. diff --git a/bpo/__init__.py b/bpo/__init__.py new file mode 100644 index 0000000..1bb6ff2 --- /dev/null +++ b/bpo/__init__.py @@ -0,0 +1 @@ +"""Behaviour Policy Optimization: an off-policy extension of PPO.""" diff --git a/bpo.py b/bpo/bpo.py similarity index 71% rename from bpo.py rename to bpo/bpo.py index 275ac55..5ac4544 100644 --- a/bpo.py +++ b/bpo/bpo.py @@ -12,19 +12,19 @@ import gymnasium as gym from gymnasium import spaces from typing import Any, Optional, TypeVar, Union, Callable -from common.base_class import BaseJaxPolicy -from common.on_policy_algorithm import OnPolicyAlgorithm -from common.policies import BPOPolicy -from common.metrics import StatsLogger, Stats, Dist -from common.buffers import BPORolloutBuffer, BPOReplayBuffer -from common.symlog import * -from ppo import PPO +from bpo.common.base_class import BaseJaxPolicy +from bpo.common.on_policy_algorithm import OnPolicyAlgorithm +from bpo.common.policies import BPOPolicy +from bpo.common.metrics import StatsLogger, Stats, Dist +from bpo.common.buffers import BPORolloutBuffer, BPOReplayBuffer +from bpo.common.symlog import * +from bpo.ppo import PPO from tqdm.auto import tqdm tfd = tfp.distributions -class BPO_PPO(PPO): +class BPO_PPO(PPO): def __init__( self, env: gym.Env, @@ -130,15 +130,15 @@ def __init__( clip_rho=self.clip_rho, clip_c=self.clip_c, clip_traj=self.clip_traj, - weight_on_adv=self.weight_on_adv + weight_on_adv=self.weight_on_adv, ) self.replay_buffer = BPOReplayBuffer( - self.replay_size//self.n_envs, + self.replay_size // self.n_envs, self.observation_space, self.action_space, self.n_envs, - gamma=self.gamma + gamma=self.gamma, ) self.replay_buffer.reset() @@ -148,43 +148,56 @@ def _setup_buffer(self): def _setup_model(self): super()._setup_model() - self.q_vf = self.policy.q_vf # type: ignore[assignment] - self.q_hat_vf = self.policy.q_hat_vf # type: ignore[assignment] + self.q_vf = self.policy.q_vf # type: ignore[assignment] + self.q_hat_vf = self.policy.q_hat_vf # type: ignore[assignment] @staticmethod - @partial(jit, static_argnames=["weight_td", "polyak_tau", "symlog_targets", "clip_targets"]) + @partial( + jit, + static_argnames=["weight_td", "polyak_tau", "symlog_targets", "clip_targets"], + ) def _update_q_vf_fqe_disc( - key: jax.Array, - featurizer_state: TrainState, - actor_state: TrainState, - q_vf_state: TrainState, - observations: jnp.ndarray, - actions: jnp.ndarray, - rewards: jnp.ndarray, - terminal: jnp.ndarray, - truncated: jnp.ndarray, - next_observations: jnp.ndarray, - mu_log_probs: jnp.ndarray, - gamma: float, - r_max: float, - r_min: float, - symlog_reg_coef: float, - weight_td: bool = False, - polyak_tau: float = 0.0, - symlog_targets: bool = False, - clip_targets: bool = False + key: jax.Array, + featurizer_state: TrainState, + actor_state: TrainState, + q_vf_state: TrainState, + observations: jnp.ndarray, + actions: jnp.ndarray, + rewards: jnp.ndarray, + terminal: jnp.ndarray, + truncated: jnp.ndarray, + next_observations: jnp.ndarray, + mu_log_probs: jnp.ndarray, + gamma: float, + r_max: float, + r_min: float, + symlog_reg_coef: float, + weight_td: bool = False, + polyak_tau: float = 0.0, + symlog_targets: bool = False, + clip_targets: bool = False, ): features = featurizer_state.apply_fn(featurizer_state.params, observations) - next_features = featurizer_state.apply_fn(featurizer_state.params, next_observations) + next_features = featurizer_state.apply_fn( + featurizer_state.params, next_observations + ) if symlog_targets: - target_q_vf_values = jnp.take_along_axis(q_vf_state.apply_fn(q_vf_state.target_params, features), actions.reshape(-1, 1), axis=1).squeeze(-1) + target_q_vf_values = jnp.take_along_axis( + q_vf_state.apply_fn(q_vf_state.target_params, features), + actions.reshape(-1, 1), + axis=1, + ).squeeze(-1) def polyak_update(target_params, source_params, tau=0.005): - return jax.tree_util.tree_map(lambda t, s: (1 - tau) * t + tau * s, target_params, source_params) + return jax.tree_util.tree_map( + lambda t, s: (1 - tau) * t + tau * s, target_params, source_params + ) - probs = actor_state.apply_fn(actor_state.params, next_features).probs_parameter() + probs = actor_state.apply_fn( + actor_state.params, next_features + ).probs_parameter() q_vf_values = q_vf_state.apply_fn(q_vf_state.params, next_features) @@ -193,10 +206,16 @@ def polyak_update(target_params, source_params, tau=0.005): next_q_vf_values = jnp.sum(probs * q_vf_values, axis=-1) - targets = rewards + gamma * (1.0 - terminal) * next_q_vf_values + gamma * truncated * next_q_vf_values + targets = ( + rewards + + gamma * (1.0 - terminal) * next_q_vf_values + + gamma * truncated * next_q_vf_values + ) if clip_targets: - clipped_targets = jnp.clip(targets, r_min / (1 - gamma), r_max / (1 - gamma)) + clipped_targets = jnp.clip( + targets, r_min / (1 - gamma), r_max / (1 - gamma) + ) clipping_ratio = jnp.mean(clipped_targets != targets) else: clipped_targets = targets @@ -219,50 +238,69 @@ def polyak_update(target_params, source_params, tau=0.005): w_norm = ratios = jnp.ones_like(mu_log_probs) def fqe_loss(params): - q_vf_values = jnp.take_along_axis(q_vf_state.apply_fn(params, features), actions.reshape(-1, 1), axis=1).squeeze(-1) - loss = jnp.mean(w_norm * ((clipped_targets - q_vf_values)**2)) + q_vf_values = jnp.take_along_axis( + q_vf_state.apply_fn(params, features), actions.reshape(-1, 1), axis=1 + ).squeeze(-1) + loss = jnp.mean(w_norm * ((clipped_targets - q_vf_values) ** 2)) if symlog_targets: - loss += symlog_reg_coef * jnp.mean(w_norm * ((q_vf_values - target_q_vf_values)**2)) + loss += symlog_reg_coef * jnp.mean( + w_norm * ((q_vf_values - target_q_vf_values) ** 2) + ) return loss - q_vf_loss, grads = jax.value_and_grad(fqe_loss, has_aux=False)(q_vf_state.params) + q_vf_loss, grads = jax.value_and_grad(fqe_loss, has_aux=False)( + q_vf_state.params + ) q_vf_state = q_vf_state.apply_gradients(grads=grads) if polyak_tau: - new_target_params = polyak_update(q_vf_state.target_params, q_vf_state.params, tau=polyak_tau) + new_target_params = polyak_update( + q_vf_state.target_params, q_vf_state.params, tau=polyak_tau + ) q_vf_state = q_vf_state.replace(target_params=new_target_params) return q_vf_state, q_vf_loss, ratios, clipping_ratio @staticmethod - @partial(jit, static_argnames=["clip_actions", "weight_td", "polyak_tau", "symlog_targets", "clip_targets"]) + @partial( + jit, + static_argnames=[ + "clip_actions", + "weight_td", + "polyak_tau", + "symlog_targets", + "clip_targets", + ], + ) def _update_q_vf_fqe_cont( - key: jax.Array, - featurizer_state: TrainState, - actor_state: TrainState, - q_vf_state: TrainState, - observations: jnp.ndarray, - actions: jnp.ndarray, - rewards: jnp.ndarray, - terminal: jnp.ndarray, - truncated: jnp.ndarray, - next_observations: jnp.ndarray, - mu_log_probs: jnp.ndarray, - clip_act_low: float, - clip_act_high: float, - gamma: float, - r_max: float, - r_min: float, - symlog_reg_coef: float, + key: jax.Array, + featurizer_state: TrainState, + actor_state: TrainState, + q_vf_state: TrainState, + observations: jnp.ndarray, + actions: jnp.ndarray, + rewards: jnp.ndarray, + terminal: jnp.ndarray, + truncated: jnp.ndarray, + next_observations: jnp.ndarray, + mu_log_probs: jnp.ndarray, + clip_act_low: float, + clip_act_high: float, + gamma: float, + r_max: float, + r_min: float, + symlog_reg_coef: float, clip_actions: bool = False, - weight_td: bool = False, - polyak_tau: float = 0.0, - symlog_targets: bool = False, - clip_targets: bool = False + weight_td: bool = False, + polyak_tau: float = 0.0, + symlog_targets: bool = False, + clip_targets: bool = False, ): features = featurizer_state.apply_fn(featurizer_state.params, observations) - next_features = featurizer_state.apply_fn(featurizer_state.params, next_observations) + next_features = featurizer_state.apply_fn( + featurizer_state.params, next_observations + ) if clip_actions: clipped_actions = jnp.clip(actions, clip_act_low, clip_act_high) @@ -270,10 +308,14 @@ def _update_q_vf_fqe_cont( clipped_actions = actions if symlog_targets: - target_q_vf_values = q_vf_state.apply_fn(q_vf_state.target_params, features, clipped_actions).squeeze(-1) + target_q_vf_values = q_vf_state.apply_fn( + q_vf_state.target_params, features, clipped_actions + ).squeeze(-1) def polyak_update(target_params, source_params, tau=0.005): - return jax.tree_util.tree_map(lambda t, s: (1 - tau) * t + tau * s, target_params, source_params) + return jax.tree_util.tree_map( + lambda t, s: (1 - tau) * t + tau * s, target_params, source_params + ) dist1 = actor_state.apply_fn(actor_state.params, next_features) next_actions = dist1.sample(seed=key) @@ -284,15 +326,25 @@ def polyak_update(target_params, source_params, tau=0.005): clipped_next_actions = next_actions if symlog_targets: - next_q_vf_values = q_vf_state.apply_fn(q_vf_state.params, next_features, clipped_next_actions).squeeze(-1) + next_q_vf_values = q_vf_state.apply_fn( + q_vf_state.params, next_features, clipped_next_actions + ).squeeze(-1) next_q_vf_values = jsymexp(next_q_vf_values) else: - next_q_vf_values = q_vf_state.apply_fn(q_vf_state.target_params, next_features, clipped_next_actions).squeeze(-1) - - targets = rewards + gamma * (1.0 - terminal) * next_q_vf_values + gamma * truncated * next_q_vf_values + next_q_vf_values = q_vf_state.apply_fn( + q_vf_state.target_params, next_features, clipped_next_actions + ).squeeze(-1) + + targets = ( + rewards + + gamma * (1.0 - terminal) * next_q_vf_values + + gamma * truncated * next_q_vf_values + ) if clip_targets: - clipped_targets = jnp.clip(targets, r_min / (1 - gamma), r_max / (1 - gamma)) + clipped_targets = jnp.clip( + targets, r_min / (1 - gamma), r_max / (1 - gamma) + ) clipping_ratio = jnp.mean(clipped_targets != targets) else: clipped_targets = targets @@ -313,17 +365,25 @@ def polyak_update(target_params, source_params, tau=0.005): w_norm = ratios = jnp.ones_like(mu_log_probs) def fqe_loss(params): - q_vf_values = q_vf_state.apply_fn(params, features, clipped_actions).squeeze(-1) - loss = jnp.mean(w_norm * ((clipped_targets - q_vf_values)**2)) + q_vf_values = q_vf_state.apply_fn( + params, features, clipped_actions + ).squeeze(-1) + loss = jnp.mean(w_norm * ((clipped_targets - q_vf_values) ** 2)) if symlog_targets: - loss += symlog_reg_coef * jnp.mean(w_norm * ((q_vf_values - target_q_vf_values)**2)) + loss += symlog_reg_coef * jnp.mean( + w_norm * ((q_vf_values - target_q_vf_values) ** 2) + ) return loss - q_vf_loss, grads = jax.value_and_grad(fqe_loss, has_aux=False)(q_vf_state.params) + q_vf_loss, grads = jax.value_and_grad(fqe_loss, has_aux=False)( + q_vf_state.params + ) q_vf_state = q_vf_state.apply_gradients(grads=grads) if polyak_tau: - new_target_params = polyak_update(q_vf_state.target_params, q_vf_state.params, tau=polyak_tau) + new_target_params = polyak_update( + q_vf_state.target_params, q_vf_state.params, tau=polyak_tau + ) q_vf_state = q_vf_state.replace(target_params=new_target_params) return q_vf_state, q_vf_loss, ratios, clipping_ratio @@ -336,17 +396,17 @@ def _copy_q_vf_params(q_vf_state): @staticmethod @partial(jit, static_argnames=["symlog_targets", "clip_targets"]) def _update_mu_disc( - key: jax.Array, - featurizer_state: TrainState, - mu_state: TrainState, - actor_state: TrainState, - q_vf_state: TrainState, - observations: jnp.ndarray, - gamma: float, - r_max: float, - ent_coef: float, - symlog_targets: bool = False, - clip_targets: bool = False + key: jax.Array, + featurizer_state: TrainState, + mu_state: TrainState, + actor_state: TrainState, + q_vf_state: TrainState, + observations: jnp.ndarray, + gamma: float, + r_max: float, + ent_coef: float, + symlog_targets: bool = False, + clip_targets: bool = False, ): features = featurizer_state.apply_fn(featurizer_state.params, observations) @@ -355,7 +415,7 @@ def mu_loss(params): dist1 = mu_state.apply_fn(params, features) entropy = dist1.entropy() - + dist2 = actor_state.apply_fn(actor_state.params, features) log_pi = jax.nn.log_softmax(dist2.logits_parameter(), axis=-1) @@ -380,27 +440,29 @@ def mu_loss(params): total_policy_loss = policy_loss + ent_coef * entropy_loss return total_policy_loss, q_hat_values - (bhv_loss, q_hat_values), grads = jax.value_and_grad(mu_loss, has_aux=True)(mu_state.params) + (bhv_loss, q_hat_values), grads = jax.value_and_grad(mu_loss, has_aux=True)( + mu_state.params + ) mu_state = mu_state.apply_gradients(grads=grads) return mu_state, bhv_loss, q_hat_values @staticmethod - @partial(jit, static_argnames=[ "mu_policy_loss", "symlog_targets", "clip_targets"]) + @partial(jit, static_argnames=["mu_policy_loss", "symlog_targets", "clip_targets"]) def _update_mu_cont( - key: jax.Array, - featurizer_state: TrainState, - mu_state: TrainState, - actor_state: TrainState, - q_vf_state: TrainState, - observations: jnp.ndarray, - gamma: float, - r_max: float, - ent_coef: float, - clip_act_low: float, - clip_act_high: float, - mu_policy_loss: str = 'kl_qhat', - symlog_targets: bool = False, - clip_targets: bool = False + key: jax.Array, + featurizer_state: TrainState, + mu_state: TrainState, + actor_state: TrainState, + q_vf_state: TrainState, + observations: jnp.ndarray, + gamma: float, + r_max: float, + ent_coef: float, + clip_act_low: float, + clip_act_high: float, + mu_policy_loss: str = "kl_qhat", + symlog_targets: bool = False, + clip_targets: bool = False, ): features = featurizer_state.apply_fn(featurizer_state.params, observations) @@ -416,10 +478,14 @@ def mu_loss(params): entropy = dist1.entropy() if symlog_targets: - q_hat_values = q_vf_state.apply_fn(q_vf_state.params, features, actions).squeeze(-1) + q_hat_values = q_vf_state.apply_fn( + q_vf_state.params, features, actions + ).squeeze(-1) q_hat_values = jsymexp(q_hat_values) else: - q_hat_values = q_vf_state.apply_fn(q_vf_state.target_params, features, actions).squeeze(-1) + q_hat_values = q_vf_state.apply_fn( + q_vf_state.target_params, features, actions + ).squeeze(-1) if clip_targets: q_hat_values = jnp.clip(q_hat_values, 1e-6, r_max / (1 - gamma)) @@ -428,10 +494,12 @@ def mu_loss(params): q_hat_values = jnp.min(q_hat_values, axis=0) - if mu_policy_loss == 'kl': + if mu_policy_loss == "kl": policy_loss = jnp.mean(mu_log_probs - actor_log_probs) - elif mu_policy_loss == 'kl_qhat': - policy_loss = jnp.mean(mu_log_probs - actor_log_probs - 0.5 * jnp.log(q_hat_values)) + elif mu_policy_loss == "kl_qhat": + policy_loss = jnp.mean( + mu_log_probs - actor_log_probs - 0.5 * jnp.log(q_hat_values) + ) else: raise NotImplementedError @@ -440,15 +508,14 @@ def mu_loss(params): total_policy_loss = policy_loss + ent_coef * entropy_loss return total_policy_loss, q_hat_values - (bhv_loss, q_hat_values), grads = jax.value_and_grad(mu_loss, has_aux=True)(mu_state.params) + (bhv_loss, q_hat_values), grads = jax.value_and_grad(mu_loss, has_aux=True)( + mu_state.params + ) mu_state = mu_state.apply_gradients(grads=grads) return mu_state, bhv_loss, q_hat_values def optimize( - self, - step: int, - logger: Optional[StatsLogger] = None, - tqdm_position: int = 1 + self, step: int, logger: Optional[StatsLogger] = None, tqdm_position: int = 1 ): ratios_stats = Stats() @@ -479,26 +546,41 @@ def optimize( self.key, key1, key2, key3, key4 = jr.split(self.key, 5) with tqdm( - total=self.n_epochs*self.n_steps//(self.batch_size//self.n_envs), + total=self.n_epochs * self.n_steps // (self.batch_size // self.n_envs), desc="optimize", position=tqdm_position, leave=False, dynamic_ncols=True, colour="cyan", ) as pbar: - for _ in range(self.n_epochs): key1, subkey = jr.split(key1) - for rollout_data in self.rollout_buffer.get(subkey, self.batch_size//self.n_envs): - - observations, actions, rewards, values, returns, advantages, mu_log_probs, actor_log_probs = rollout_data + for rollout_data in self.rollout_buffer.get( + subkey, self.batch_size // self.n_envs + ): + ( + observations, + actions, + rewards, + values, + returns, + advantages, + mu_log_probs, + actor_log_probs, + ) = rollout_data if isinstance(self.action_space, spaces.Discrete): # Convert discrete action from float to int actions = actions.flatten().astype(np.int32) - (self.policy.featurizer_state, self.policy.actor_state, self.policy.critic_state), (pg_loss, vf_loss) = \ - self._one_update( + ( + ( + self.policy.featurizer_state, + self.policy.actor_state, + self.policy.critic_state, + ), + (pg_loss, vf_loss), + ) = self._one_update( featurizer_state=self.policy.featurizer_state, actor_state=self.policy.actor_state, critic_state=self.policy.critic_state, @@ -506,7 +588,9 @@ def optimize( actions=actions, advantages=advantages, returns=returns, - old_log_prob=mu_log_probs if not self.weight_on_adv else actor_log_probs, + old_log_prob=mu_log_probs + if not self.weight_on_adv + else actor_log_probs, clip_range=float(clip_range), ent_coef=0.0, vf_coef=self.vf_coef, @@ -531,11 +615,26 @@ def optimize( for _ in range(self.n_vf_epochs): key2, subkey = jr.split(key2) - for replay_data in self.replay_buffer.get(subkey, self.vf_batch_size//self.n_envs): - - observations, actions, rewards, _, terminal, truncated, next_observations, mu_log_probs = replay_data + for replay_data in self.replay_buffer.get( + subkey, self.vf_batch_size // self.n_envs + ): + ( + observations, + actions, + rewards, + _, + terminal, + truncated, + next_observations, + mu_log_probs, + ) = replay_data - assert np.all(np.logical_or(np.logical_not(truncated.astype(np.float32) == 1.0), (terminal.astype(np.float32) == 1.0))) + assert np.all( + np.logical_or( + np.logical_not(truncated.astype(np.float32) == 1.0), + (terminal.astype(np.float32) == 1.0), + ) + ) if isinstance(self.action_space, spaces.Discrete): # Convert discrete action from float to int @@ -544,8 +643,12 @@ def optimize( key2, subkey = jr.split(key2) if isinstance(self.action_space, spaces.Discrete): - self.policy.q_vf_state, q_vf_loss, ratios_2, q_vf_clipping_ratios = \ - self._update_q_vf_fqe_disc( + ( + self.policy.q_vf_state, + q_vf_loss, + ratios_2, + q_vf_clipping_ratios, + ) = self._update_q_vf_fqe_disc( key=subkey, featurizer_state=self.policy.featurizer_state, actor_state=self.policy.actor_state, @@ -567,8 +670,12 @@ def optimize( clip_targets=self.clip_targets, ) elif isinstance(self.action_space, spaces.Box): - self.policy.q_vf_state, q_vf_loss, ratios_2, q_vf_clipping_ratios = \ - self._update_q_vf_fqe_cont( + ( + self.policy.q_vf_state, + q_vf_loss, + ratios_2, + q_vf_clipping_ratios, + ) = self._update_q_vf_fqe_cont( key=subkey, featurizer_state=self.policy.featurizer_state, actor_state=self.policy.actor_state, @@ -602,18 +709,17 @@ def optimize( ratios_2_dist.update(np.array(ratios_2)) q_vf_clipping_ratio_stats.update(np.array(q_vf_clipping_ratios)) - if not self.polyak_tau: self.policy.q_vf_state = self._copy_q_vf_params(self.policy.q_vf_state) q_values, r_hat_rewards = self.replay_buffer.compute_r_hat_rewards( self.policy.featurizer_state, - self.policy.q_vf_state, + self.policy.q_vf_state, clip_actions=self.clip_actions, r_max=self.r_max, r_min=self.r_min, - symlog_targets=self.symlog_targets, - clip_targets=self.clip_targets + symlog_targets=self.symlog_targets, + clip_targets=self.clip_targets, ) if self.verbose > 0: @@ -624,9 +730,19 @@ def optimize( for _ in range(self.n_vf_epochs): key3, subkey = jr.split(key3) - for replay_data in self.replay_buffer.get(subkey, self.vf_batch_size//self.n_envs): - - observations, actions, _, r_hat_rewards, terminal, truncated, next_observations, mu_log_probs = replay_data + for replay_data in self.replay_buffer.get( + subkey, self.vf_batch_size // self.n_envs + ): + ( + observations, + actions, + _, + r_hat_rewards, + terminal, + truncated, + next_observations, + mu_log_probs, + ) = replay_data if isinstance(self.action_space, spaces.Discrete): # Convert discrete action from float to int @@ -635,8 +751,12 @@ def optimize( key3, subkey = jr.split(key3) if isinstance(self.action_space, spaces.Discrete): - self.policy.q_hat_vf_state, q_hat_vf_loss, _, q_hat_vf_clipping_ratios = \ - self._update_q_vf_fqe_disc( + ( + self.policy.q_hat_vf_state, + q_hat_vf_loss, + _, + q_hat_vf_clipping_ratios, + ) = self._update_q_vf_fqe_disc( key=subkey, featurizer_state=self.policy.featurizer_state, actor_state=self.policy.actor_state, @@ -658,8 +778,12 @@ def optimize( clip_targets=self.clip_targets, ) elif isinstance(self.action_space, spaces.Box): - self.policy.q_hat_vf_state, q_hat_vf_loss, _, q_hat_vf_clipping_ratios = \ - self._update_q_vf_fqe_cont( + ( + self.policy.q_hat_vf_state, + q_hat_vf_loss, + _, + q_hat_vf_clipping_ratios, + ) = self._update_q_vf_fqe_cont( key=subkey, featurizer_state=self.policy.featurizer_state, actor_state=self.policy.actor_state, @@ -688,23 +812,26 @@ def optimize( if self.verbose > 0: q_hat_vf_loss_stats.update(float(q_hat_vf_loss)) - q_hat_vf_clipping_ratio_stats.update(np.array(q_hat_vf_clipping_ratios)) - + q_hat_vf_clipping_ratio_stats.update( + np.array(q_hat_vf_clipping_ratios) + ) if not self.polyak_tau: - self.policy.q_hat_vf_state = self._copy_q_vf_params(self.policy.q_hat_vf_state) + self.policy.q_hat_vf_state = self._copy_q_vf_params( + self.policy.q_hat_vf_state + ) for _ in range(self.n_mu_epochs): key4, subkey = jr.split(key4) - for replay_data in self.replay_buffer.get(subkey, self.mu_batch_size//self.n_envs): - + for replay_data in self.replay_buffer.get( + subkey, self.mu_batch_size // self.n_envs + ): observations, _, _, _, _, _, _, _ = replay_data key4, subkey = jr.split(key4) if isinstance(self.action_space, spaces.Discrete): - self.policy.mu_state, mu_loss, q_hat_values = \ - self._update_mu_disc( + self.policy.mu_state, mu_loss, q_hat_values = self._update_mu_disc( key=subkey, featurizer_state=self.policy.featurizer_state, mu_state=self.policy.mu_state, @@ -718,8 +845,7 @@ def optimize( clip_targets=self.clip_targets, ) elif isinstance(self.action_space, spaces.Box): - self.policy.mu_state, mu_loss, q_hat_values = \ - self._update_mu_cont( + self.policy.mu_state, mu_loss, q_hat_values = self._update_mu_cont( key=subkey, featurizer_state=self.policy.featurizer_state, mu_state=self.policy.mu_state, @@ -793,7 +919,7 @@ def optimize( logger.add(logs) def rollout( - self, + self, step: int, logger: Optional[RolloutLogger] = None, tqdm_position: int = 1, @@ -816,27 +942,35 @@ def rollout( dynamic_ncols=True, colour="green", ) - if self.use_tqdm_rollout else nullcontext() + if self.use_tqdm_rollout + else nullcontext() ) with pbar_context as pbar: while steps < self.n_steps: - - if self.use_sde and self.sde_sample_freq > 0 and steps % self.sde_sample_freq == 0: + if ( + self.use_sde + and self.sde_sample_freq > 0 + and steps % self.sde_sample_freq == 0 + ): self.policy.reset_noise() if not self.use_sde or isinstance(self.action_space, spaces.Discrete): self.policy.reset_noise() obs = self.prepare_obs(self._last_obs, n_envs=self.n_envs) - actions, mu_log_probs, actor_log_probs, values = self.policy.predict_all(self.policy.noise_key, obs) + actions, mu_log_probs, actor_log_probs, values = ( + self.policy.predict_all(self.policy.noise_key, obs) + ) actions = np.array(actions) mu_log_probs = np.array(mu_log_probs) actor_log_probs = np.array(actor_log_probs) values = np.array(values) - new_obs, rewards, terminated, truncated, infos = self.env.step(self.prepare_act(actions, n_envs=self.n_envs)) + new_obs, rewards, terminated, truncated, infos = self.env.step( + self.prepare_act(actions, n_envs=self.n_envs) + ) new_obs = np.array(new_obs) rewards = np.array(rewards) @@ -845,7 +979,7 @@ def rollout( self.r_min = min(self.r_min, np.min(rewards)) steps += 1 - + if self.use_tqdm_rollout: pbar.update(1) @@ -854,14 +988,16 @@ def rollout( actions = np.array(actions) actions = actions.reshape(-1, 1) - dones = np.array([False]*self.n_envs) + dones = np.array([False] * self.n_envs) bootstrap_rewards = rewards.copy() for idx, info in enumerate(infos): if truncated[idx]: truncated_obs = new_obs[idx].reshape(1, -1) - feats = self.featurizer.apply(self.policy.featurizer_state.params, truncated_obs) + feats = self.featurizer.apply( + self.policy.featurizer_state.params, truncated_obs + ) terminal_value = np.array( self.critic.apply( self.policy.critic_state.params, @@ -889,7 +1025,7 @@ def rollout( self.replay_buffer.add( self._last_obs, actions, - rewards, # don't use bootstrapped rewards here we want to use the q functions to bootstrap + rewards, # don't use bootstrapped rewards here we want to use the q functions to bootstrap dones, truncated, next_obs, @@ -908,8 +1044,8 @@ def rollout( if logger: for info in infos: logger.add(info) - - assert isinstance(self._last_obs, np.ndarray) + + assert isinstance(self._last_obs, np.ndarray) final_obs = self.prepare_obs(self._last_obs, n_envs=self.n_envs) feats = self.featurizer.apply(self.policy.featurizer_state.params, final_obs) last_value = np.array( @@ -919,11 +1055,10 @@ def rollout( ).flatten() ) - self.rollout_buffer.compute_returns_and_advantages(last_value=last_value, done=self._last_episode_start) + self.rollout_buffer.compute_returns_and_advantages( + last_value=last_value, done=self._last_episode_start + ) @property def train_ratio(self): return self.n_steps * self.n_envs - - - diff --git a/bpo/common/__init__.py b/bpo/common/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/common/base_class.py b/bpo/common/base_class.py similarity index 78% rename from common/base_class.py rename to bpo/common/base_class.py index 4a2a644..2b96627 100644 --- a/common/base_class.py +++ b/bpo/common/base_class.py @@ -1,17 +1,17 @@ from __future__ import annotations from typing import Any, Optional, TypeVar, Union, List, Callable -from common.wrappers import NormWrapper +from bpo.common.wrappers import NormWrapper from gymnasium import spaces import gymnasium as gym import jax.random as jr -from common.metrics import RolloutLogger, StatsLogger +from bpo.common.metrics import RolloutLogger, StatsLogger from abc import ABC, abstractmethod import tensorflow as tf import math from tqdm.auto import tqdm -class BaseAlgorithm(ABC): +class BaseAlgorithm(ABC): def __init__( self, env: gym.Env, @@ -22,7 +22,7 @@ def __init__( verbose: int = 0, supported_action_spaces: Optional[tuple[type[spaces.Space], ...]] = None, supported_observation_spaces: Optional[tuple[type[spaces.Space], ...]] = None, - eval_env: Optional[gym.Env] = None, + eval_env: Optional[gym.Env] = None, ): self.env = env @@ -33,7 +33,7 @@ def __init__( self.supported_action_spaces = supported_action_spaces self.supported_observation_spaces = supported_observation_spaces self.eval_env = eval_env - + act_sp = getattr(env, "action_space", None) obs_sp = getattr(env, "observation_space", None) @@ -42,16 +42,20 @@ def __init__( f"{self.__class__.__name__}: env is missing action_space/observation_space." ) - if self.supported_action_spaces is not None and not isinstance(act_sp, self.supported_action_spaces): + if self.supported_action_spaces is not None and not isinstance( + act_sp, self.supported_action_spaces + ): raise TypeError( - f"{self.__class__.__name__} does not support action space { type(act_sp).__name__ }.\n" - f"Allowed types: { _allowed_names(self.supported_action_spaces) }." + f"{self.__class__.__name__} does not support action space {type(act_sp).__name__}.\n" + f"Allowed types: {_allowed_names(self.supported_action_spaces)}." ) - if self.supported_observation_spaces is not None and not isinstance(obs_sp, self.supported_observation_spaces): + if self.supported_observation_spaces is not None and not isinstance( + obs_sp, self.supported_observation_spaces + ): raise TypeError( - f"{self.__class__.__name__} does not support observation space { type(obs_sp).__name__ }.\n" - f"Allowed types: { _allowed_names(self.supported_observation_spaces) }." + f"{self.__class__.__name__} does not support observation space {type(obs_sp).__name__}.\n" + f"Allowed types: {_allowed_names(self.supported_observation_spaces)}." ) self.observation_space = env.observation_space @@ -75,28 +79,28 @@ def train( summary_writer = None rollout_logger = RolloutLogger( - stdout=True, - tqdm=True, - tensorboard=bool(summary_writer is not None), - summary_writer=summary_writer, + stdout=True, + tqdm=True, + tensorboard=bool(summary_writer is not None), + summary_writer=summary_writer, stats_window_size=stats_window_size, - prefix="train/rollout" + prefix="train/rollout", ) eval_logger = RolloutLogger( - stdout=True, - tqdm=True, - tensorboard=bool(summary_writer is not None), - summary_writer=summary_writer, + stdout=True, + tqdm=True, + tensorboard=bool(summary_writer is not None), + summary_writer=summary_writer, stats_window_size=num_eval_episodes, - prefix="eval/rollout" + prefix="eval/rollout", ) stats_logger = StatsLogger( - stdout=True, - tqdm=True, - tensorboard=bool(summary_writer is not None), - summary_writer=summary_writer, + stdout=True, + tqdm=True, + tensorboard=bool(summary_writer is not None), + summary_writer=summary_writer, stats_window_size=stats_window_size, - prefix="train/stats" + prefix="train/stats", ) total_steps = 0 @@ -112,8 +116,7 @@ def train( leave=True, dynamic_ncols=True, ) as iter_bar: - - for iteration in range(math.ceil((num_frames)/self.train_ratio)): + for iteration in range(math.ceil((num_frames) / self.train_ratio)): self.rollout(total_steps, logger=rollout_logger) self.optimize(total_steps, logger=stats_logger) @@ -133,9 +136,7 @@ def train( def _get_eval_env(self) -> gym.Env: if self.eval_env is not None: return self.eval_env - raise RuntimeError( - "Cannot construct eval env; please pass env_fn or eval_env." - ) + raise RuntimeError("Cannot construct eval env; please pass env_fn or eval_env.") def optimize(self, step: int, logger: Optional[StatsLogger] = None): """Optimizes the policy and other auxilliary models""" @@ -145,7 +146,12 @@ def rollout(self, step: int, logger: Optional[RolloutLogger] = None): """Collects rollouts/experience for the policy to learn from""" raise NotImplementedError - def eval(self, num_episodes: int, seed: Optional[int] = None, logger: Optional[RolloutLogger] = None) -> List[float]: + def eval( + self, + num_episodes: int, + seed: Optional[int] = None, + logger: Optional[RolloutLogger] = None, + ) -> List[float]: eval_env = self._get_eval_env() base = 0 if self.seed is None else int(self.seed) eval_seed = base + 72 if seed is None else int(seed) + 72 @@ -157,9 +163,8 @@ def eval(self, num_episodes: int, seed: Optional[int] = None, logger: Optional[R position=1, leave=False, dynamic_ncols=True, - colour="yellow" + colour="yellow", ) as pbar: - for ep in range(num_episodes): obs, info = self.eval_env.reset(seed=eval_seed + ep) done = False @@ -188,12 +193,12 @@ def train_ratio(self) -> int: """How often to do an update step during training""" raise NotImplementedError -class BaseJaxPolicy(ABC): +class BaseJaxPolicy(ABC): def __init__( self, observation_space: spaces.Space, action_space: spaces.Space, ): self.observation_space = observation_space - self.action_space = action_space \ No newline at end of file + self.action_space = action_space diff --git a/common/buffers.py b/bpo/common/buffers.py similarity index 63% rename from common/buffers.py rename to bpo/common/buffers.py index 48403d8..666821f 100644 --- a/common/buffers.py +++ b/bpo/common/buffers.py @@ -1,4 +1,3 @@ - import gymnasium as gym from gymnasium import spaces import numpy as np @@ -9,17 +8,17 @@ import jax.numpy as jnp from jax import jit from functools import partial -from common.symlog import * +from bpo.common.symlog import * -class RolloutBuffer(): +class RolloutBuffer: def __init__( - self, - buffer_size: int, - observation_space: spaces.Space, - action_space: spaces.Space, - n_envs: int, - gamma: float = 0.99, + self, + buffer_size: int, + observation_space: spaces.Space, + action_space: spaces.Space, + n_envs: int, + gamma: float = 0.99, gae_lambda: float = 1.0, ): @@ -29,7 +28,7 @@ def __init__( self.n_envs = n_envs self.gamma = gamma self.gae_lambda = gae_lambda - + self.obs_shape = self.observation_space.shape if isinstance(self.action_space, spaces.Box): @@ -44,18 +43,28 @@ def __init__( raise NotImplementedError def reset(self): - self.observations = np.zeros((self.buffer_size, self.n_envs, *self.obs_shape), dtype=np.float32) - self.actions = np.zeros((self.buffer_size, self.n_envs, self.act_dim), dtype=np.float32) + self.observations = np.zeros( + (self.buffer_size, self.n_envs, *self.obs_shape), dtype=np.float32 + ) + self.actions = np.zeros( + (self.buffer_size, self.n_envs, self.act_dim), dtype=np.float32 + ) self.rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) self.returns = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) - self.episode_starts = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) + self.episode_starts = np.zeros( + (self.buffer_size, self.n_envs), dtype=np.float32 + ) self.values = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) self.log_probs = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) self.advantages = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) self.pos = 0 self.full = False - def compute_returns_and_advantages(self, last_value: Optional[np.ndarray] = None, done: np.ndarray = np.array(False)): + def compute_returns_and_advantages( + self, + last_value: Optional[np.ndarray] = None, + done: np.ndarray = np.array(False), + ): if last_value is None: last_value = np.array(0.0) @@ -67,22 +76,30 @@ def compute_returns_and_advantages(self, last_value: Optional[np.ndarray] = None next_non_terminal = 1.0 - done.astype(np.float32) next_value = last_value else: - next_non_terminal = 1.0 - self.episode_starts[step + 1].astype(np.float32) - next_value = self.values[step +1] - delta = self.rewards[step] + self.gamma * next_value * next_non_terminal - self.values[step] - last_gae_lam = delta + self.gamma * self.gae_lambda * next_non_terminal * last_gae_lam + next_non_terminal = 1.0 - self.episode_starts[step + 1].astype( + np.float32 + ) + next_value = self.values[step + 1] + delta = ( + self.rewards[step] + + self.gamma * next_value * next_non_terminal + - self.values[step] + ) + last_gae_lam = ( + delta + self.gamma * self.gae_lambda * next_non_terminal * last_gae_lam + ) self.advantages[step] = last_gae_lam self.returns = self.advantages + self.values def add( - self, - obs: np.ndarray, + self, + obs: np.ndarray, actions: np.ndarray, - rewards: np.ndarray, + rewards: np.ndarray, episode_starts: np.ndarray, - values: np.ndarray, - log_probs: np.ndarray + values: np.ndarray, + log_probs: np.ndarray, ): if self.full: @@ -104,7 +121,7 @@ def add( self.pos += 1 if self.pos == self.buffer_size: self.full = True - + return True def get(self, key: jax.Array, batch_size: int): @@ -112,32 +129,34 @@ def get(self, key: jax.Array, batch_size: int): indices = jr.permutation(key, self.buffer_size) if batch_size is None: - batch_size=self.buffer_size + batch_size = self.buffer_size start_idx = 0 while start_idx < self.buffer_size: - yield self._get_samples(indices[start_idx: start_idx + batch_size]) + yield self._get_samples(indices[start_idx : start_idx + batch_size]) start_idx += batch_size def _get_samples(self, batch_inds: np.ndarray): - data = (self.observations[batch_inds].reshape(-1, *self.obs_shape), - self.actions[batch_inds].reshape(-1, self.act_dim), - self.rewards[batch_inds].flatten(), - self.values[batch_inds].flatten(), - self.returns[batch_inds].flatten(), - self.advantages[batch_inds].flatten(), - self.log_probs[batch_inds].flatten()) + data = ( + self.observations[batch_inds].reshape(-1, *self.obs_shape), + self.actions[batch_inds].reshape(-1, self.act_dim), + self.rewards[batch_inds].flatten(), + self.values[batch_inds].flatten(), + self.returns[batch_inds].flatten(), + self.advantages[batch_inds].flatten(), + self.log_probs[batch_inds].flatten(), + ) return data -class BPOReplayBuffer(): +class BPOReplayBuffer: def __init__( - self, - buffer_size: int, - observation_space: spaces.Space, + self, + buffer_size: int, + observation_space: spaces.Space, action_space: spaces.Space, - n_envs: int, - gamma: float = 0.99 + n_envs: int, + gamma: float = 0.99, ): self.buffer_size = buffer_size @@ -163,13 +182,19 @@ def __init__( raise NotImplementedError def reset(self): - self.observations = np.zeros((self.buffer_size, self.n_envs, *self.obs_shape), dtype=np.float32) - self.actions = np.zeros((self.buffer_size, self.n_envs, self.act_dim), dtype=np.float32) + self.observations = np.zeros( + (self.buffer_size, self.n_envs, *self.obs_shape), dtype=np.float32 + ) + self.actions = np.zeros( + (self.buffer_size, self.n_envs, self.act_dim), dtype=np.float32 + ) self.rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) self.r_hat_rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) self.terminal = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) self.truncated = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) - self.next_observations = np.zeros((self.buffer_size, self.n_envs, *self.obs_shape), dtype=np.float32) + self.next_observations = np.zeros( + (self.buffer_size, self.n_envs, *self.obs_shape), dtype=np.float32 + ) self.mu_log_probs = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) self.pos = 0 self.full = False @@ -177,47 +202,51 @@ def reset(self): @staticmethod @partial(jit, static_argnames=["discrete_acts", "symlog_targets", "clip_targets"]) def _one_q_vf_bootstrap( - featurizer_state: TrainState, - q_vf_state: TrainState, - obs: jnp.ndarray, - act: jnp.ndarray, - discrete_acts: bool, - gamma: float, - r_max: float, - r_min: float, - symlog_targets: bool, - clip_targets: bool + featurizer_state: TrainState, + q_vf_state: TrainState, + obs: jnp.ndarray, + act: jnp.ndarray, + discrete_acts: bool, + gamma: float, + r_max: float, + r_min: float, + symlog_targets: bool, + clip_targets: bool, ): feats = featurizer_state.apply_fn(featurizer_state.params, obs) if discrete_acts: - q_value = jnp.take_along_axis(q_vf_state.apply_fn(q_vf_state.params, feats), act, axis=1).squeeze(-1) + q_value = jnp.take_along_axis( + q_vf_state.apply_fn(q_vf_state.params, feats), act, axis=1 + ).squeeze(-1) if symlog_targets: q_value = jsymexp(q_value) if clip_targets: q_value = jnp.clip(q_value, r_min / (1 - gamma), r_max / (1 - gamma)) return q_value - + if symlog_targets: q_value = q_vf_state.apply_fn(q_vf_state.params, feats, act).squeeze(-1) q_value = jsymexp(q_value) else: - q_value = q_vf_state.apply_fn(q_vf_state.target_params, feats, act).squeeze(-1) + q_value = q_vf_state.apply_fn(q_vf_state.target_params, feats, act).squeeze( + -1 + ) if clip_targets: q_value = jnp.clip(q_value, r_min / (1 - gamma), r_max / (1 - gamma)) - + return q_value def compute_r_hat_rewards( - self, - featurizer_state: TrainState, - q_vf_state: TrainState, - clip_actions: bool = False, - r_max: float = 1.0, - r_min: float = 0.0, - symlog_targets: bool = False, - clip_targets: bool = False + self, + featurizer_state: TrainState, + q_vf_state: TrainState, + clip_actions: bool = False, + r_max: float = 1.0, + r_min: float = 0.0, + symlog_targets: bool = False, + clip_targets: bool = False, ): q_values = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) @@ -237,34 +266,36 @@ def compute_r_hat_rewards( else: actions = jnp.asarray(self.actions[step], dtype=jnp.float32) if clip_actions: - actions = jnp.clip(actions, self.action_space.low, self.action_space.high) + actions = jnp.clip( + actions, self.action_space.low, self.action_space.high + ) q_values[step] = self._one_q_vf_bootstrap( featurizer_state=featurizer_state, - q_vf_state=q_vf_state, + q_vf_state=q_vf_state, obs=observations, act=actions, - discrete_acts=discrete_acts, + discrete_acts=discrete_acts, gamma=jnp.float32(self.gamma), r_max=jnp.float32(r_max), r_min=jnp.float32(r_min), - symlog_targets=symlog_targets, + symlog_targets=symlog_targets, clip_targets=clip_targets, ) - self.r_hat_rewards = 2 * self.rewards * q_values - self.rewards ** 2 + self.r_hat_rewards = 2 * self.rewards * q_values - self.rewards**2 return q_values, self.r_hat_rewards def add( - self, - obs: np.ndarray, - actions: np.ndarray, - rewards: np.ndarray, - terminal: np.ndarray, - truncated: np.ndarray, - next_obs: np.ndarray, - mu_log_probs: np.ndarray + self, + obs: np.ndarray, + actions: np.ndarray, + rewards: np.ndarray, + terminal: np.ndarray, + truncated: np.ndarray, + next_obs: np.ndarray, + mu_log_probs: np.ndarray, ): if len(mu_log_probs.shape) == 0: @@ -285,7 +316,7 @@ def add( if self.pos == self.buffer_size: self.full = True self.pos = 0 - + return True def get(self, key: jax.Array, batch_size: int): @@ -293,11 +324,11 @@ def get(self, key: jax.Array, batch_size: int): indices = jr.permutation(key, self.buffer_size) if batch_size is None: - batch_size=self.buffer_size + batch_size = self.buffer_size start_idx = 0 while start_idx < self.buffer_size: - yield self._get_samples(indices[start_idx: start_idx + batch_size]) + yield self._get_samples(indices[start_idx : start_idx + batch_size]) start_idx += batch_size else: indices = jr.permutation(key, self.pos) @@ -310,34 +341,36 @@ def get(self, key: jax.Array, batch_size: int): start_idx = 0 while start_idx < self.pos: - yield self._get_samples(indices[start_idx: start_idx + batch_size]) + yield self._get_samples(indices[start_idx : start_idx + batch_size]) start_idx += batch_size def _get_samples(self, batch_inds: np.ndarray): - data = (self.observations[batch_inds].reshape(-1, *self.obs_shape), - self.actions[batch_inds].reshape(-1, self.act_dim), - self.rewards[batch_inds].flatten(), - self.r_hat_rewards[batch_inds].flatten(), - self.terminal[batch_inds].flatten(), - self.truncated[batch_inds].flatten(), - self.next_observations[batch_inds].reshape(-1, *self.obs_shape), - self.mu_log_probs[batch_inds].flatten(),) + data = ( + self.observations[batch_inds].reshape(-1, *self.obs_shape), + self.actions[batch_inds].reshape(-1, self.act_dim), + self.rewards[batch_inds].flatten(), + self.r_hat_rewards[batch_inds].flatten(), + self.terminal[batch_inds].flatten(), + self.truncated[batch_inds].flatten(), + self.next_observations[batch_inds].reshape(-1, *self.obs_shape), + self.mu_log_probs[batch_inds].flatten(), + ) return data -class BPORolloutBuffer(): +class BPORolloutBuffer: def __init__( - self, - buffer_size: int, - observation_space: spaces.Space, - action_space: spaces.Space, - n_envs: int, - gamma: float = 0.99, - gae_lambda: float = 1.0, - clip_rho: float = 1.0, + self, + buffer_size: int, + observation_space: spaces.Space, + action_space: spaces.Space, + n_envs: int, + gamma: float = 0.99, + gae_lambda: float = 1.0, + clip_rho: float = 1.0, clip_c: float = 1.0, clip_traj: bool = False, - weight_on_adv: bool = False + weight_on_adv: bool = False, ): self.buffer_size = buffer_size @@ -350,7 +383,7 @@ def __init__( self.clip_c = clip_c self.clip_traj = clip_traj self.weight_on_adv = weight_on_adv - + self.pos = 0 self.full = False @@ -368,19 +401,31 @@ def __init__( raise NotImplementedError def reset(self): - self.observations = np.zeros((self.buffer_size, self.n_envs, *self.obs_shape), dtype=np.float32) - self.actions = np.zeros((self.buffer_size, self.n_envs, self.act_dim), dtype=np.float32) + self.observations = np.zeros( + (self.buffer_size, self.n_envs, *self.obs_shape), dtype=np.float32 + ) + self.actions = np.zeros( + (self.buffer_size, self.n_envs, self.act_dim), dtype=np.float32 + ) self.rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) self.returns = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) - self.episode_starts = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) + self.episode_starts = np.zeros( + (self.buffer_size, self.n_envs), dtype=np.float32 + ) self.values = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) self.mu_log_probs = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) - self.actor_log_probs = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) + self.actor_log_probs = np.zeros( + (self.buffer_size, self.n_envs), dtype=np.float32 + ) self.advantages = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32) self.pos = 0 self.full = False - def compute_returns_and_advantages(self, last_value: Optional[np.ndarray] = None, done: np.ndarray = np.array(False)): + def compute_returns_and_advantages( + self, + last_value: Optional[np.ndarray] = None, + done: np.ndarray = np.array(False), + ): # Remember that importance-weighted estimators can exhibit high variance. Consider applying clipping/truncation strategies (such as those in Retrace or V-trace) for improved stability in practice. # See IMPALA/V-Trace clipping coefficients c @@ -399,10 +444,16 @@ def compute_returns_and_advantages(self, last_value: Optional[np.ndarray] = None next_non_terminal = 1.0 - done.astype(np.float32) next_value = last_value else: - next_non_terminal = 1.0 - self.episode_starts[step + 1].astype(np.float32) - next_value = self.values[step +1] - - delta = self.rewards[step] + self.gamma * next_value * next_non_terminal - self.values[step] + next_non_terminal = 1.0 - self.episode_starts[step + 1].astype( + np.float32 + ) + next_value = self.values[step + 1] + + delta = ( + self.rewards[step] + + self.gamma * next_value * next_non_terminal + - self.values[step] + ) deltas[step] = delta @@ -410,14 +461,22 @@ def compute_returns_and_advantages(self, last_value: Optional[np.ndarray] = None rho = np.minimum(ratios[step], self.clip_rho) c = np.minimum(traj_c[step:], self.clip_c) - decay_weights = (self.gamma * self.gae_lambda) ** (np.arange(step, self.buffer_size) - step) + decay_weights = (self.gamma * self.gae_lambda) ** ( + np.arange(step, self.buffer_size) - step + ) - self.returns[step] = self.values[step] + rho * np.sum(c * deltas[step:] * decay_weights[:, np.newaxis], axis=0) + self.returns[step] = self.values[step] + rho * np.sum( + c * deltas[step:] * decay_weights[:, np.newaxis], axis=0 + ) if self.weight_on_adv: - self.advantages[step] = rho * np.sum(c * deltas[step:] * decay_weights[:, np.newaxis], axis=0) + self.advantages[step] = rho * np.sum( + c * deltas[step:] * decay_weights[:, np.newaxis], axis=0 + ) else: - self.advantages[step] = np.sum(c * deltas[step:] * decay_weights[:, np.newaxis], axis=0) + self.advantages[step] = np.sum( + c * deltas[step:] * decay_weights[:, np.newaxis], axis=0 + ) # accumulate ratios traj_c[step:] *= ratios[step][np.newaxis, :] @@ -426,7 +485,9 @@ def compute_returns_and_advantages(self, last_value: Optional[np.ndarray] = None rho = np.minimum(ratios[step], self.clip_rho) c = np.minimum(ratios[step], self.clip_c) - bootstrap = next_non_terminal * ((1 - self.gae_lambda) * next_value + self.gae_lambda * last_ret_lam) + bootstrap = next_non_terminal * ( + (1 - self.gae_lambda) * next_value + self.gae_lambda * last_ret_lam + ) gae = self.rewards[step] - self.values[step] + self.gamma * bootstrap @@ -435,18 +496,26 @@ def compute_returns_and_advantages(self, last_value: Optional[np.ndarray] = None else: self.advantages[step] = gae - last_ret_lam = self.values[step] + rho * delta + self.gamma * self.gae_lambda * next_non_terminal * c * (last_ret_lam - next_value) + last_ret_lam = ( + self.values[step] + + rho * delta + + self.gamma + * self.gae_lambda + * next_non_terminal + * c + * (last_ret_lam - next_value) + ) self.returns[step] = last_ret_lam def add( - self, - obs: np.ndarray, - actions: np.ndarray, - rewards: np.ndarray, - episode_starts: np.ndarray, - values: np.ndarray, - mu_log_probs: np.ndarray, - actor_log_probs: np.ndarray + self, + obs: np.ndarray, + actions: np.ndarray, + rewards: np.ndarray, + episode_starts: np.ndarray, + values: np.ndarray, + mu_log_probs: np.ndarray, + actor_log_probs: np.ndarray, ): if self.full: @@ -473,7 +542,7 @@ def add( self.pos += 1 if self.pos == self.buffer_size: self.full = True - + return True def get(self, key: jax.Array, batch_size: int): @@ -481,20 +550,22 @@ def get(self, key: jax.Array, batch_size: int): indices = jr.permutation(key, self.buffer_size) if batch_size is None: - batch_size=self.buffer_size + batch_size = self.buffer_size start_idx = 0 while start_idx < self.buffer_size: - yield self._get_samples(indices[start_idx: start_idx + batch_size]) + yield self._get_samples(indices[start_idx : start_idx + batch_size]) start_idx += batch_size def _get_samples(self, batch_inds: np.ndarray): - data = (self.observations[batch_inds].reshape(-1, *self.obs_shape), - self.actions[batch_inds].reshape(-1, self.act_dim), - self.rewards[batch_inds].flatten(), - self.values[batch_inds].flatten(), - self.returns[batch_inds].flatten(), - self.advantages[batch_inds].flatten(), - self.mu_log_probs[batch_inds].flatten(), - self.actor_log_probs[batch_inds].flatten()) - return data \ No newline at end of file + data = ( + self.observations[batch_inds].reshape(-1, *self.obs_shape), + self.actions[batch_inds].reshape(-1, self.act_dim), + self.rewards[batch_inds].flatten(), + self.values[batch_inds].flatten(), + self.returns[batch_inds].flatten(), + self.advantages[batch_inds].flatten(), + self.mu_log_probs[batch_inds].flatten(), + self.actor_log_probs[batch_inds].flatten(), + ) + return data diff --git a/bpo/common/configs.py b/bpo/common/configs.py new file mode 100644 index 0000000..927f1cb --- /dev/null +++ b/bpo/common/configs.py @@ -0,0 +1,514 @@ +import re +import sys +import io +import json + +import contextlib +import glob +import os +import re +import shutil + + +class Path: + filesystems = [] + + def __new__(cls, path): + path = str(path) + for impl, pred in cls.filesystems: + if pred(path): + obj = super().__new__(impl) + obj.__init__(path) + return obj + raise NotImplementedError(f"No filesystem supports: {path}") + + def __getnewargs__(self): + return (self._path,) + + def __init__(self, path): + assert isinstance(path, str) + path = re.sub(r"^\./*", "", path) # Remove leading dot or dot slashes. + path = re.sub(r"(?<=[^/])/$", "", path) # Remove single trailing slash. + path = path or "." # Empty path is represented by a dot. + self._path = path + + def __truediv__(self, part): + sep = "" if self._path.endswith("/") else "/" + return type(self)(f"{self._path}{sep}{str(part)}") + + def __repr__(self): + return f"Path({str(self)})" + + def __fspath__(self): + return str(self) + + def __eq__(self, other): + return self._path == other._path + + def __lt__(self, other): + return self._path < other._path + + def __str__(self): + return self._path + + @property + def parent(self): + if "/" not in self._path: + return type(self)(".") + parent = self._path.rsplit("/", 1)[0] + parent = parent or ("/" if self._path.startswith("/") else ".") + return type(self)(parent) + + @property + def name(self): + if "/" not in self._path: + return self._path + return self._path.rsplit("/", 1)[1] + + @property + def stem(self): + return self.name.split(".", 1)[0] if "." in self.name else self.name + + @property + def suffix(self): + return ("." + self.name.split(".", 1)[1]) if "." in self.name else "" + + def read(self, mode="r"): + assert mode in "r rb".split(), mode + with self.open(mode) as f: + return f.read() + + def write(self, content, mode="w"): + assert mode in "w a wb ab".split(), mode + with self.open(mode) as f: + f.write(content) + + @contextlib.contextmanager + def open(self, mode="r"): + raise NotImplementedError + + def absolute(self): + raise NotImplementedError + + def glob(self, pattern): + raise NotImplementedError + + def exists(self): + raise NotImplementedError + + def isfile(self): + raise NotImplementedError + + def isdir(self): + raise NotImplementedError + + def mkdirs(self): + raise NotImplementedError + + def remove(self): + raise NotImplementedError + + def rmtree(self): + raise NotImplementedError + + def copy(self, dest): + raise NotImplementedError + + def move(self, dest): + self.copy(dest) + self.remove() + + +class LocalPath(Path): + def __init__(self, path): + super().__init__(os.path.expanduser(str(path))) + + @contextlib.contextmanager + def open(self, mode="r"): + with open(str(self), mode=mode) as f: + yield f + + def absolute(self): + return type(self)(os.path.absolute(str(self))) + + def glob(self, pattern): + for path in glob.glob(f"{str(self)}/{pattern}"): + yield type(self)(path) + + def exists(self): + return os.path.exists(str(self)) + + def isfile(self): + return os.path.isfile(str(self)) + + def isdir(self): + return os.path.isdir(str(self)) + + def mkdirs(self): + os.makedirs(str(self), exist_ok=True) + + def remove(self): + os.rmdir(str(self)) if self.isdir() else os.remove(str(self)) + + def rmtree(self): + shutil.rmtree(self) + + def copy(self, dest): + if self.isfile(): + shutil.copy(self, type(self)(dest)) + else: + shutil.copytree(self, type(self)(dest), dirs_exist_ok=True) + + def move(self, dest): + shutil.move(self, dest) + + +class GFilePath(Path): + def __init__(self, path): + path = str(path) + if not (path.startswith("/") or "://" in path): + path = os.path.abspath(os.path.expanduser(path)) + super().__init__(path) + import tensorflow as tf + + self._gfile = tf.io.gfile + + @contextlib.contextmanager + def open(self, mode="r"): + path = str(self) + if "a" in mode and path.startswith("/cns/"): + path += "%r=3.2" + if mode.startswith("x") and self.exists(): + raise FileExistsError(path) + mode = mode.replace("x", "w") + with self._gfile.GFile(path, mode) as f: + yield f + + def absolute(self): + return self + + def glob(self, pattern): + for path in self._gfile.glob(f"{str(self)}/{pattern}"): + yield type(self)(path) + + def exists(self): + return self._gfile.exists(str(self)) + + def isfile(self): + return self.exists() and not self.isdir() + + def isdir(self): + return self._gfile.isdir(str(self)) + + def mkdirs(self): + self._gfile.makedirs(str(self)) + + def remove(self): + self._gfile.remove(str(self)) + + def rmtree(self): + self._gfile.rmtree(str(self)) + + def copy(self, dest): + self._gfile.copy(str(self), str(dest), overwrite=True) + + def move(self, dest): + dest = Path(dest) + if dest.isdir(): + dest.rmtree() + self._gfile.rename(self, str(dest), overwrite=True) + + +Path.filesystems = [ + (GFilePath, lambda path: path.startswith("gs://")), + (GFilePath, lambda path: path.startswith("/cns/")), + (LocalPath, lambda path: True), +] + + +class Flags: + def __init__(self, *args, **kwargs): + self._config = Config(*args, **kwargs) + + def parse(self, argv=None, help_exists=True): + parsed, remaining = self.parse_known(argv) + for flag in remaining: + if flag.startswith("--"): + raise ValueError(f"Flag '{flag}' did not match any config keys.") + # assert not remaining or , remaining + return parsed + + def parse_known(self, argv=None, help_exists=False): + if argv is None: + argv = sys.argv[1:] + if "--help" in argv: + print("\nHelp:") + lines = str(self._config).split("\n")[2:] + print("\n".join("--" + re.sub(r"[:,\[\]]", "", x) for x in lines)) + help_exists and sys.exit() + parsed = {} + remaining = [] + key = None + vals = None + for arg in argv: + if arg.startswith("--"): + if key: + self._submit_entry(key, vals, parsed, remaining) + if "=" in arg: + key, val = arg.split("=", 1) + vals = [val] + else: + key, vals = arg, [] + else: + if key: + vals.append(arg) + else: + remaining.append(arg) + self._submit_entry(key, vals, parsed, remaining) + parsed = self._config.update(parsed) + return parsed, remaining + + def _submit_entry(self, key, vals, parsed, remaining): + if not key and not vals: + return + if not key: + vals = ", ".join(f"'{x}'" for x in vals) + raise ValueError(f"Values {vals} were not preceded by any flag.") + name = key[len("--") :] + if "=" in name: + remaining.extend([key] + vals) + return + if self._config.IS_PATTERN.fullmatch(name): + pattern = re.compile(name) + keys = {k for k in self._config.flat if pattern.fullmatch(k)} + elif name in self._config: + keys = [name] + else: + keys = [] + if not keys: + remaining.extend([key] + vals) + return + if not vals: + raise ValueError(f"Flag '{key}' was not followed by any values.") + for key in keys: + parsed[key] = self._parse_flag_value(self._config[key], vals, key) + + def _parse_flag_value(self, default, value, key): + value = value if isinstance(value, (tuple, list)) else (value,) + if isinstance(default, (tuple, list)): + if len(value) == 1 and "," in value[0]: + value = value[0].split(",") + return tuple(self._parse_flag_value(default[0], [x], key) for x in value) + assert len(value) == 1, value + value = str(value[0]) + if default is None: + return value + if isinstance(default, bool): + try: + return bool(["False", "True"].index(value)) + except ValueError: + message = f"Expected bool but got '{value}' for key '{key}'." + raise TypeError(message) + if isinstance(default, int): + try: + value = float(value) # Allow scientific notation for integers. + assert float(int(value)) == value + except (TypeError, AssertionError): + message = f"Expected int but got float '{value}' for key '{key}'." + raise TypeError(message) + return int(value) + if isinstance(default, dict): + raise TypeError( + f"Key '{key}' refers to a whole dict. Please speicfy a subkey." + ) + return type(default)(value) + + +class Config(dict): + SEP = "." + IS_PATTERN = re.compile(r".*[^A-Za-z0-9_.-].*") + + def __init__(self, *args, **kwargs): + mapping = dict(*args, **kwargs) + mapping = self._flatten(mapping) + mapping = self._ensure_keys(mapping) + mapping = self._ensure_values(mapping) + self._flat = mapping + self._nested = self._nest(mapping) + # Need to assign the values to the base class dictionary so that + # conversion to dict does not lose the content. + super().__init__(self._nested) + + @property + def flat(self): + return self._flat.copy() + + def save(self, filename): + filename = path.Path(filename) + if filename.suffix == ".json": + filename.write(json.dumps(dict(self))) + elif filename.suffix in (".yml", ".yaml"): + import ruamel.yaml as yaml + + with io.StringIO() as stream: + yaml.safe_dump(dict(self), stream) + filename.write(stream.getvalue()) + else: + raise NotImplementedError(filename.suffix) + + @classmethod + def load(cls, filename): + filename = path.Path(filename) + if filename.suffix == ".json": + return cls(json.loads(filename.read_text())) + elif filename.suffix in (".yml", ".yaml"): + import ruamel.yaml as yaml + + return cls(yaml.safe_load(filename.read_text())) + else: + raise NotImplementedError(filename.suffix) + + def __contains__(self, name): + try: + self[name] + return True + except KeyError: + return False + + def __getattr__(self, name): + if name.startswith("_"): + return super().__getattr__(name) + try: + return self[name] + except KeyError: + raise AttributeError(name) + + def __getitem__(self, name): + result = self._nested + for part in name.split(self.SEP): + try: + result = result[part] + except TypeError: + raise KeyError + if isinstance(result, dict): + result = type(self)(result) + return result + + def __setattr__(self, key, value): + if key.startswith("_"): + return super().__setattr__(key, value) + message = f"Tried to set key '{key}' on immutable config. Use update()." + raise AttributeError(message) + + def __setitem__(self, key, value): + if key.startswith("_"): + return super().__setitem__(key, value) + message = f"Tried to set key '{key}' on immutable config. Use update()." + raise AttributeError(message) + + def __reduce__(self): + return (type(self), (dict(self),)) + + def __str__(self): + lines = ["\nConfig:"] + keys, vals, typs = [], [], [] + for key, val in self.flat.items(): + keys.append(key + ":") + vals.append(self._format_value(val)) + typs.append(self._format_type(val)) + max_key = max(len(k) for k in keys) if keys else 0 + max_val = max(len(v) for v in vals) if vals else 0 + for key, val, typ in zip(keys, vals, typs): + key = key.ljust(max_key) + val = val.ljust(max_val) + lines.append(f"{key} {val} ({typ})") + return "\n".join(lines) + + def update(self, *args, **kwargs): + result = self._flat.copy() + inputs = self._flatten(dict(*args, **kwargs)) + for key, new in inputs.items(): + if self.IS_PATTERN.match(key): + pattern = re.compile(key) + keys = {k for k in result if pattern.match(k)} + else: + keys = [key] + if not keys: + raise KeyError(f"Unknown key or pattern {key}.") + for key in keys: + old = result[key] + try: + if isinstance(old, int) and isinstance(new, float): + if float(int(new)) != new: + message = f"Cannot convert fractional float {new} to int." + raise ValueError(message) + result[key] = type(old)(new) + except (ValueError, TypeError): + raise TypeError( + f"Cannot convert '{new}' to type '{type(old).__name__}' " + + f"for key '{key}' with previous value '{old}'." + ) + return type(self)(result) + + def _flatten(self, mapping): + result = {} + for key, value in mapping.items(): + if isinstance(value, dict): + for k, v in self._flatten(value).items(): + if self.IS_PATTERN.match(key) or self.IS_PATTERN.match(k): + combined = f"{key}\\{self.SEP}{k}" + else: + combined = f"{key}{self.SEP}{k}" + result[combined] = v + else: + result[key] = value + return result + + def _nest(self, mapping): + result = {} + for key, value in mapping.items(): + parts = key.split(self.SEP) + node = result + for part in parts[:-1]: + if part not in node: + node[part] = {} + node = node[part] + node[parts[-1]] = value + return result + + def _ensure_keys(self, mapping): + for key in mapping: + assert not self.IS_PATTERN.match(key), key + return mapping + + def _ensure_values(self, mapping): + result = json.loads(json.dumps(mapping)) + for key, value in result.items(): + if isinstance(value, list): + value = list(value) + if isinstance(value, tuple): + if len(value) == 0: + message = ( + "Empty lists are disallowed because their type is unclear." + ) + raise TypeError(message) + if not isinstance(value[0], (str, float, int, bool)): + message = "Lists can only contain strings, floats, ints, bools" + message += f" but not {type(value[0])}" + raise TypeError(message) + if not all(isinstance(x, type(value[0])) for x in value[1:]): + message = "Elements of a list must all be of the same type." + raise TypeError(message) + result[key] = value + return result + + def _format_value(self, value): + if isinstance(value, (list, tuple)): + return "[" + ", ".join(self._format_value(x) for x in value) + "]" + return str(value) + + def _format_type(self, value): + if isinstance(value, (list, tuple)): + assert len(value) > 0, value + return self._format_type(value[0]) + "s" + return str(type(value).__name__) diff --git a/common/layers.py b/bpo/common/layers.py similarity index 66% rename from common/layers.py rename to bpo/common/layers.py index 55c056c..4a319f1 100644 --- a/common/layers.py +++ b/bpo/common/layers.py @@ -8,6 +8,7 @@ Dtype = Any Array = Any + class Flatten(nn.Module): """ Equivalent to PyTorch nn.Flatten() layer. @@ -17,12 +18,13 @@ class Flatten(nn.Module): def __call__(self, x: jnp.ndarray) -> jnp.ndarray: return x.reshape((x.shape[0], -1)) -class Identity(nn.Module): +class Identity(nn.Module): @nn.compact def __call__(self, x: jnp.ndarray) -> jnp.ndarray: return x + class NatureCNN(nn.Module): grayscale_obs: bool = True normalize_images: bool = True @@ -37,10 +39,28 @@ def __call__(self, x: jnp.ndarray) -> jnp.ndarray: x = jnp.transpose(x, (0, 4, 2, 3, 1)) if self.normalize_images: x = x.astype(jnp.float32) / 255.0 - x = nn.Conv(32, kernel_size=(8, 8), strides=(4, 4), padding="VALID", kernel_init=self.kernel_init)(x) + x = nn.Conv( + 32, + kernel_size=(8, 8), + strides=(4, 4), + padding="VALID", + kernel_init=self.kernel_init, + )(x) x = self.activation_fn(x) - x = nn.Conv(64, kernel_size=(4, 4), strides=(2, 2), padding="VALID", kernel_init=self.kernel_init)(x) + x = nn.Conv( + 64, + kernel_size=(4, 4), + strides=(2, 2), + padding="VALID", + kernel_init=self.kernel_init, + )(x) x = self.activation_fn(x) - x = nn.Conv(64, kernel_size=(3, 3), strides=(1, 1), padding="VALID", kernel_init=self.kernel_init)(x) + x = nn.Conv( + 64, + kernel_size=(3, 3), + strides=(1, 1), + padding="VALID", + kernel_init=self.kernel_init, + )(x) x = self.activation_fn(x) - return x.reshape(x.shape[0], -1) \ No newline at end of file + return x.reshape(x.shape[0], -1) diff --git a/common/metrics.py b/bpo/common/metrics.py similarity index 93% rename from common/metrics.py rename to bpo/common/metrics.py index 267c1a3..1b48584 100644 --- a/common/metrics.py +++ b/bpo/common/metrics.py @@ -19,6 +19,7 @@ import tensorflow as tf from tqdm import tqdm + class Stats: r"""Streaming scalar summary statistics for an arbitrary batch of values. @@ -148,9 +149,13 @@ def __add__(self, other: "Stats") -> "Stats": RuntimeError: If either object has not observed any data. """ assert isinstance(other, Stats) - assert self.prefix == other.prefix, "can't add two Stats objects with different prefixes" + assert self.prefix == other.prefix, ( + "can't add two Stats objects with different prefixes" + ) if self.stats is None or other.stats is None: - raise RuntimeError("Cannot add Stats objects before both have been updated at least once.") + raise RuntimeError( + "Cannot add Stats objects before both have been updated at least once." + ) new = Stats(prefix=self.prefix) new.n = int(self.n + other.n) @@ -158,7 +163,8 @@ def __add__(self, other: "Stats") -> "Stats": w_self = self.n / new.n w_other = other.n / new.n new.stats = { - "mean": float(self.stats["mean"]) * w_self + float(other.stats["mean"]) * w_other, + "mean": float(self.stats["mean"]) * w_self + + float(other.stats["mean"]) * w_other, "mean_squares": float(self.stats["mean_squares"]) * w_self + float(other.stats["mean_squares"]) * w_other, "max": float(max(float(self.stats["max"]), float(other.stats["max"]))), @@ -167,6 +173,7 @@ def __add__(self, other: "Stats") -> "Stats": } return new + class Dist: r"""Reservoir-sampled distribution summary. @@ -254,10 +261,12 @@ def __add__(self, other: "Dist") -> "Dist": AssertionError: If :attr:`prefix` differs. """ assert isinstance(other, Dist) - assert self.prefix == other.prefix, "can't add two Dist objects with different prefixes" - assert ( - self.reservoir_size == other.reservoir_size - ), "can't add two Dist objects with different reservoir sizes" + assert self.prefix == other.prefix, ( + "can't add two Dist objects with different prefixes" + ) + assert self.reservoir_size == other.reservoir_size, ( + "can't add two Dist objects with different reservoir sizes" + ) new = Dist(prefix=self.prefix, reservoir_size=self.reservoir_size, rng=0) new.n = int(self.n + other.n) @@ -272,6 +281,7 @@ def __add__(self, other: "Dist") -> "Dist": new.res = both[idx].astype(np.float32, copy=False) return new + class BaseLogger: r"""Base class for logging scalar statistics and distributions. @@ -317,7 +327,9 @@ def __init__( self.summary_writer: Optional[tf.summary.SummaryWriter] = summary_writer if self.tensorboard: - assert self.summary_writer is not None, "tensorboard=True requires a summary_writer" + assert self.summary_writer is not None, ( + "tensorboard=True requires a summary_writer" + ) if (self.summary_writer is not None) and (not self.tensorboard): warnings.warn( "tensorboard is set to False but summary writer is provided; this may produce unexpected behaviour", @@ -325,7 +337,9 @@ def __init__( ) self.stats_window_size: int = int(stats_window_size) - self.prefix: str = (prefix if not prefix else (prefix if prefix.endswith("/") else prefix + "/")) + self.prefix: str = ( + prefix if not prefix else (prefix if prefix.endswith("/") else prefix + "/") + ) self.stats: Dict[str, Deque[float]] = {} def reset(self): @@ -356,6 +370,7 @@ def log(self, step: int): """ raise NotImplementedError + class StatsLogger(BaseLogger): r"""Logger for streaming scalar stats (:class:`Stats`) and distributions (:class:`Dist`). @@ -396,7 +411,6 @@ def __init__( self.stats_to_log: Dict[str, float] = {} self.dists_to_log: Dict[str, np.ndarray] = {} - def reset(self): """Clear all buffered scalar and distribution values.""" super().reset() @@ -421,7 +435,9 @@ def add(self, new: Mapping[str, Union["Stats", "Dist", float, int, np.floating]] if f"{key}_{k}" in self.stats: self.stats[f"{key}_{k}"].append(float(v)) else: - self.stats[f"{key}_{k}"] = deque([float(v)], maxlen=self.stats_window_size) + self.stats[f"{key}_{k}"] = deque( + [float(v)], maxlen=self.stats_window_size + ) elif isinstance(val, Dist): # Store a snapshot of the reservoir for later histogram logging. self.dists[key] = val.get() @@ -487,16 +503,20 @@ def _log_to_stdout(self, step: int): Args: step: Global step index (unused; included for API symmetry). """ - stats_to_print = {key: "{0:.4g}".format(val) for key, val in self.stats_to_log.items()} + stats_to_print = { + key: "{0:.4g}".format(val) for key, val in self.stats_to_log.items() + } if not stats_to_print: return - max_key_len = max([len(key) for key in stats_to_print] + [max(0, len(self.prefix) - 2)]) + max_key_len = max( + [len(key) for key in stats_to_print] + [max(0, len(self.prefix) - 2)] + ) max_val_len = max([len(val) for val in stats_to_print.values()]) stdout = "" max_len = 1 + 4 + max_key_len + 2 + 1 + 2 + max_val_len + 2 + 1 - stdout += ("-" * max_len + "\n") + stdout += "-" * max_len + "\n" stdout += ( "| " + self.prefix @@ -515,13 +535,14 @@ def _log_to_stdout(self, step: int): + " " * (max_val_len - len(val) + 2) + "|\n" ) - stdout += ("-" * max_len + "\n") + stdout += "-" * max_len + "\n" if self.tqdm: tqdm.write(stdout) else: print(stdout) + class RolloutLogger(BaseLogger): r"""Logger for episodic metrics produced during environment rollouts. @@ -611,7 +632,7 @@ def _add_scalars(self, scalars: Mapping[str, Union[float, int, np.floating]]): else: # +1 because we keep the most recent value as "current episode". self.stats[k] = deque([float(v)], maxlen=self.stats_window_size + 1) - + def _create_stats_to_log(self): """Compute per-metric mean over *completed* episodes.""" self.stats_to_log = {} @@ -619,7 +640,9 @@ def _create_stats_to_log(self): if len(val) > 1: # Temporarily drop last (current) value from the mean. last = val.pop() - self.stats_to_log[key] = float(np.mean(val)) if len(val) > 0 else float(last) + self.stats_to_log[key] = ( + float(np.mean(val)) if len(val) > 0 else float(last) + ) val.append(last) def _log_to_tensorboard(self, step: int): @@ -655,12 +678,14 @@ def _log_to_stdout(self, step: int): if not stats_to_print: return - max_key_len = max([len(key) for key in stats_to_print] + [max(0, len(self.prefix) - 2)]) + max_key_len = max( + [len(key) for key in stats_to_print] + [max(0, len(self.prefix) - 2)] + ) max_val_len = max([len(val) for val in stats_to_print.values()]) stdout = "" max_len = 1 + 4 + max_key_len + 2 + 1 + 2 + max_val_len + 2 + 1 - stdout += ("-" * max_len + "\n") + stdout += "-" * max_len + "\n" stdout += ( "| " + self.prefix @@ -679,17 +704,9 @@ def _log_to_stdout(self, step: int): + " " * (max_val_len - len(val) + 2) + "|\n" ) - stdout += ("-" * max_len + "\n") + stdout += "-" * max_len + "\n" if self.tqdm: tqdm.write(stdout) else: print(stdout) - - - - - - - - diff --git a/common/on_policy_algorithm.py b/bpo/common/on_policy_algorithm.py similarity index 73% rename from common/on_policy_algorithm.py rename to bpo/common/on_policy_algorithm.py index 66d7962..69efaa1 100644 --- a/common/on_policy_algorithm.py +++ b/bpo/common/on_policy_algorithm.py @@ -9,21 +9,30 @@ import numpy as np import gymnasium as gym from gymnasium import spaces -from common.base_class import BaseAlgorithm, BaseJaxPolicy +from bpo.common.base_class import BaseAlgorithm, BaseJaxPolicy from tqdm import tqdm from typing import Any, Optional, TypeVar, Union, Callable -from common.wrappers import DummyVecWrapper, VecEnvWrapperBase, NormWrapper, VecNormWrapper, OneHotObsWrapper, FlattenDictObsWrapper, is_wrapped -from common.buffers import RolloutBuffer -from common.policies import PPOPolicy -from common.metrics import RolloutLogger +from bpo.common.wrappers import ( + DummyVecWrapper, + VecEnvWrapperBase, + NormWrapper, + VecNormWrapper, + OneHotObsWrapper, + FlattenDictObsWrapper, + is_wrapped, +) +from bpo.common.buffers import RolloutBuffer +from bpo.common.policies import PPOPolicy +from bpo.common.metrics import RolloutLogger from tqdm.auto import tqdm + @contextmanager def nullcontext(arg=None): yield arg -class OnPolicyAlgorithm(BaseAlgorithm): +class OnPolicyAlgorithm(BaseAlgorithm): def __init__( self, env: gym.Env, @@ -32,7 +41,7 @@ def __init__( monitor: bool = True, device: str = "auto", verbose: int = 0, - eval_env: Optional[gym.Env] = None, + eval_env: Optional[gym.Env] = None, use_tqdm_rollout: bool = False, learning_rate: Union[float, optax.Schedule] = 3e-4, n_steps: int = 16, @@ -57,7 +66,12 @@ def __init__( monitor=monitor, device=device, verbose=verbose, - supported_action_spaces=(spaces.Discrete, spaces.Box, spaces.MultiBinary, spaces.MultiDiscrete), + supported_action_spaces=( + spaces.Discrete, + spaces.Box, + spaces.MultiBinary, + spaces.MultiDiscrete, + ), supported_observation_spaces=(spaces.Box), eval_env=eval_env, ) @@ -70,7 +84,9 @@ def __init__( if isinstance(learning_rate, float): self.lr_schedule = optax.schedules.constant_schedule(learning_rate) else: - assert callable(learning_rate), f"learning_rate for class PPO must be float or optax.Schedule not {type(learning_rate)}" + assert callable(learning_rate), ( + f"learning_rate for class PPO must be float or optax.Schedule not {type(learning_rate)}" + ) self.lr_schedule = learning_rate self.use_tqdm_rollout = use_tqdm_rollout @@ -92,35 +108,37 @@ def __init__( def _setup_buffer(self): self.rollout_buffer = RolloutBuffer( - self.n_steps, - self.observation_space, - self.action_space, + self.n_steps, + self.observation_space, + self.action_space, self.n_envs, gamma=self.gamma, - gae_lambda=self.gae_lambda + gae_lambda=self.gae_lambda, ) def _setup_model(self): if not hasattr(self, "policy") or self.policy is None: self.policy = self.policy_class( - self.observation_space, - self.action_space, - **self.policy_kwargs + self.observation_space, self.action_space, **self.policy_kwargs ) self.key = self.policy.build(self.key, self.lr_schedule, self.max_grad_norm) - self.featurizer = self.policy.featurizer # type: ignore[assignment] + self.featurizer = self.policy.featurizer # type: ignore[assignment] self.actor = self.policy.actor # type: ignore[assignment] self.critic = self.policy.critic # type: ignore[assignment] def _wrap_env(self, env: gym.Env) -> gym.Env: if is_wrapped(env, VecEnvWrapperBase): - assert is_wrapped(env, OneHotObsWrapper) or not isinstance(env.observation_space, spaces.Discrete), \ - "Before wrapping your environment in a DummyVecWrapper / VecWrapper, please wrap your environment in a OneHotObsWrapper if it has type(observation_space)__name__ = Discrete" - assert not isinstance(env.observation_space, spaces.Dict), \ - "Before wrapping your environment in a DummyVecWrapper / VecWrapper, please wrap your environment in a FlattenDictObsWrapper if it has type(observation_space)__name__ = Dict" + assert is_wrapped(env, OneHotObsWrapper) or not isinstance( + env.observation_space, spaces.Discrete + ), ( + "Before wrapping your environment in a DummyVecWrapper / VecWrapper, please wrap your environment in a OneHotObsWrapper if it has type(observation_space)__name__ = Discrete" + ) + assert not isinstance(env.observation_space, spaces.Dict), ( + "Before wrapping your environment in a DummyVecWrapper / VecWrapper, please wrap your environment in a FlattenDictObsWrapper if it has type(observation_space)__name__ = Dict" + ) if isinstance(env.observation_space, (spaces.Dict, spaces.Discrete)): if not is_wrapped(env, OneHotObsWrapper): @@ -128,7 +146,9 @@ def _wrap_env(self, env: gym.Env) -> gym.Env: if isinstance(env.observation_space, spaces.Dict): env = FlattenDictObsWrapper(env) - if not isinstance(env, VecEnvWrapperBase) and not is_wrapped(env, VecEnvWrapperBase): + if not isinstance(env, VecEnvWrapperBase) and not is_wrapped( + env, VecEnvWrapperBase + ): env = DummyVecWrapper(env) return env @@ -148,18 +168,17 @@ def _get_eval_env(self) -> gym.Env: self._eval_env = FlattenDictObsWrapper(self._eval_env) if is_wrapped(self.env, (VecNormWrapper, NormWrapper)): - if not isinstance(self._eval_env, NormWrapper): self._eval_env = NormWrapper( - env=self._eval_env, # already VecEnvWrapperBase - norm_obs=self.env.norm_obs, - norm_rew=self.env.norm_rew, - training=False, # important: do not update stats during eval - clip_obs=self.env.clip_obs, - clip_rew=self.env.clip_rew, - gamma=self.env.gamma, - eps=self.env.eps, - ) + env=self._eval_env, # already VecEnvWrapperBase + norm_obs=self.env.norm_obs, + norm_rew=self.env.norm_rew, + training=False, # important: do not update stats during eval + clip_obs=self.env.clip_obs, + clip_rew=self.env.clip_rew, + gamma=self.env.gamma, + eps=self.env.eps, + ) else: self._eval_env.training = False @@ -171,11 +190,11 @@ def _get_eval_env(self) -> gym.Env: return self._eval_env def train(self, *args, **kwargs): - self._last_episode_start = [True]*self.n_envs + self._last_episode_start = [True] * self.n_envs super().train(*args, **kwargs) def rollout( - self, + self, step: int, logger: Optional[RolloutLogger] = None, tqdm_position: int = 1, @@ -197,32 +216,40 @@ def rollout( dynamic_ncols=True, colour="green", ) - if self.use_tqdm_rollout else nullcontext() + if self.use_tqdm_rollout + else nullcontext() ) with pbar_context as pbar: while steps < self.n_steps: - - if self.use_sde and self.sde_sample_freq > 0 and steps % self.sde_sample_freq == 0: + if ( + self.use_sde + and self.sde_sample_freq > 0 + and steps % self.sde_sample_freq == 0 + ): self.policy.reset_noise() if not self.use_sde or isinstance(self.action_space, spaces.Discrete): self.policy.reset_noise() obs = self.prepare_obs(self._last_obs, n_envs=self.n_envs) - actions, log_probs, values = self.policy.predict_all(self.policy.noise_key, obs) + actions, log_probs, values = self.policy.predict_all( + self.policy.noise_key, obs + ) actions = np.array(actions) log_probs = np.array(log_probs) values = np.array(values) - new_obs, rewards, terminated, truncated, infos = self.env.step(self.prepare_act(actions, n_envs=self.n_envs)) + new_obs, rewards, terminated, truncated, infos = self.env.step( + self.prepare_act(actions, n_envs=self.n_envs) + ) new_obs = np.array(new_obs) rewards = np.array(rewards) steps += 1 - + if self.use_tqdm_rollout: pbar.update(1) @@ -231,12 +258,14 @@ def rollout( actions = np.array(actions) actions = actions.reshape(-1, 1) - dones = np.array([False]*self.n_envs) + dones = np.array([False] * self.n_envs) for idx, info in enumerate(infos): if truncated[idx]: truncated_obs = new_obs[idx].reshape(1, -1) - feats = self.featurizer.apply(self.policy.featurizer_state.params, truncated_obs) + feats = self.featurizer.apply( + self.policy.featurizer_state.params, truncated_obs + ) terminal_value = np.array( self.critic.apply( self.policy.critic_state.params, @@ -270,7 +299,7 @@ def rollout( for info in infos: logger.add(info) - assert isinstance(self._last_obs, np.ndarray) + assert isinstance(self._last_obs, np.ndarray) final_obs = self.prepare_obs(self._last_obs, n_envs=self.n_envs) feats = self.featurizer.apply(self.policy.featurizer_state.params, final_obs) last_value = np.array( @@ -280,16 +309,22 @@ def rollout( ).flatten() ) - self.rollout_buffer.compute_returns_and_advantages(last_value=last_value, done=self._last_episode_start) + self.rollout_buffer.compute_returns_and_advantages( + last_value=last_value, done=self._last_episode_start + ) - def act(self, key: jax.Array, obs: np.ndarray, deterministic: bool = False) -> Union[int, np.ndarray]: + def act( + self, key: jax.Array, obs: np.ndarray, deterministic: bool = False + ) -> Union[int, np.ndarray]: obs = self.prepare_obs(obs, n_envs=1) action = self.policy.forward(key, obs, deterministic=deterministic) return self.prepare_act(action, n_envs=1) def prepare_act(self, act: np.ndarray, n_envs: int = 1) -> Union[int, np.ndarray]: if isinstance(self.action_space, spaces.Box): - act = np.clip(act, self.action_space.low, self.action_space.high, dtype=np.float32) + act = np.clip( + act, self.action_space.low, self.action_space.high, dtype=np.float32 + ) act = act.reshape(n_envs, *self.action_space.shape) if n_envs == 1: act = np.squeeze(act, axis=0) @@ -310,12 +345,11 @@ def prepare_act(self, act: np.ndarray, n_envs: int = 1) -> Union[int, np.ndarray if n_envs == 1: act = act[0] else: - raise NotImplementedError(f"Unexpected action space {type(self.action_space).__name__}") + raise NotImplementedError( + f"Unexpected action space {type(self.action_space).__name__}" + ) return act def prepare_obs(self, obs: np.ndarray, n_envs: int = 1) -> np.ndarray: obs = np.array(obs) return obs.reshape(n_envs, *self.observation_space.shape) - - - diff --git a/common/policies.py b/bpo/common/policies.py similarity index 81% rename from common/policies.py rename to bpo/common/policies.py index ff73ef7..25faf0c 100644 --- a/common/policies.py +++ b/bpo/common/policies.py @@ -1,4 +1,5 @@ import os + os.environ["TF_USE_LEGACY_KERAS"] = "1" from dataclasses import field @@ -15,9 +16,9 @@ from flax.linen import initializers import tensorflow_probability.substrates.jax as tfp from typing import Callable, Optional, Sequence, Union, Tuple, Any -from common.layers import Flatten, Identity, NatureCNN +from bpo.common.layers import Flatten, Identity, NatureCNN from gymnasium import spaces -from common.base_class import BaseJaxPolicy +from bpo.common.base_class import BaseJaxPolicy tfd = tfp.distributions @@ -26,9 +27,11 @@ Dtype = Any Array = Any + class ValueTrainState(TrainState): target_params: Any + class ContinuousCritic(nn.Module): n_units: int = 256 activation_fn: Callable[[jnp.ndarray], jnp.ndarray] = nn.tanh @@ -37,7 +40,7 @@ class ContinuousCritic(nn.Module): kernel_init: Callable[[PRNGKey, Shape, Dtype], Array] = initializers.lecun_normal() output_dim: int = 1 softplus_final: bool = False - eps: float = 1e-6 + eps: float = 1e-6 @nn.compact def __call__(self, x: jnp.ndarray, action: jnp.ndarray) -> jnp.ndarray: @@ -59,6 +62,7 @@ def __call__(self, x: jnp.ndarray, action: jnp.ndarray) -> jnp.ndarray: x = nn.softplus(x) + self.eps return x + class DiscreteCritic(nn.Module): n_units: int = 256 activation_fn: Callable[[jnp.ndarray], jnp.ndarray] = nn.tanh @@ -67,7 +71,7 @@ class DiscreteCritic(nn.Module): kernel_init: Callable[[PRNGKey, Shape, Dtype], Array] = initializers.lecun_normal() output_dim: int = 1 softplus_final: bool = False - eps: float = 1e-6 + eps: float = 1e-6 @nn.compact def __call__(self, x: jnp.ndarray) -> jnp.ndarray: @@ -88,13 +92,14 @@ def __call__(self, x: jnp.ndarray) -> jnp.ndarray: x = nn.softplus(x) + self.eps return x + class Critic(nn.Module): n_units: int = 256 activation_fn: Callable[[jnp.ndarray], jnp.ndarray] = nn.tanh use_layer_norm: bool = False kernel_init: Callable[[PRNGKey, Shape, Dtype], Array] = initializers.lecun_normal() softplus_final: bool = False - eps: float = 1e-6 + eps: float = 1e-6 @nn.compact def __call__(self, x: jnp.ndarray) -> jnp.ndarray: @@ -112,6 +117,7 @@ def __call__(self, x: jnp.ndarray) -> jnp.ndarray: x = nn.softplus(x) + self.eps return x + class Actor(nn.Module): output_dim: int n_units: int = 256 @@ -146,8 +152,12 @@ def __call__(self, x: jnp.ndarray) -> tfd.Distribution: # type: ignore[name-def action_logits = nn.Dense(self.output_dim, kernel_init=self.kernel_init)(x) if self.num_discrete_choices is None: # Continuous actions - log_std = self.param("log_std", constant(self.log_std_init), (self.output_dim,)) - dist = tfd.MultivariateNormalDiag(loc=action_logits, scale_diag=jnp.exp(log_std)) + log_std = self.param( + "log_std", constant(self.log_std_init), (self.output_dim,) + ) + dist = tfd.MultivariateNormalDiag( + loc=action_logits, scale_diag=jnp.exp(log_std) + ) elif isinstance(self.num_discrete_choices, int): dist = tfd.Categorical(logits=action_logits) else: @@ -170,12 +180,13 @@ def __call__(self, x: jnp.ndarray) -> tfd.Distribution: # type: ignore[name-def axis=1, ) dist = tfp.distributions.Independent( - tfp.distributions.Categorical(logits=logits_padded), reinterpreted_batch_ndims=1 + tfp.distributions.Categorical(logits=logits_padded), + reinterpreted_batch_ndims=1, ) return dist -class PPOPolicy(BaseJaxPolicy): +class PPOPolicy(BaseJaxPolicy): def __init__( self, observation_space: spaces.Space, @@ -192,10 +203,7 @@ def __init__( **kwargs, ): - super().__init__( - observation_space, - action_space - ) + super().__init__(observation_space, action_space) if optimizer_kwargs is None: optimizer_kwargs = {} @@ -223,7 +231,12 @@ def __init__( self.key = self.noise_key = jax.random.PRNGKey(0) - def build(self, key: jax.Array, lr_schedule: Union[optax.Schedule, float], max_grad_norm: float): + def build( + self, + key: jax.Array, + lr_schedule: Union[optax.Schedule, float], + max_grad_norm: float, + ): key, feat_key, actor_key, critic_key = jax.random.split(key, 4) key, self.key = jax.random.split(key, 2) @@ -232,43 +245,49 @@ def build(self, key: jax.Array, lr_schedule: Union[optax.Schedule, float], max_g obs = jnp.array([self.observation_space.sample()]) if isinstance(self.action_space, spaces.Box): - self.actor_kwargs.update({ - "output_dim": int(np.prod(self.action_space.shape)), - }) + self.actor_kwargs.update( + { + "output_dim": int(np.prod(self.action_space.shape)), + } + ) elif isinstance(self.action_space, spaces.Discrete): - self.actor_kwargs.update({ - "output_dim": int(self.action_space.n), - "num_discrete_choices": int(self.action_space.n), - }) + self.actor_kwargs.update( + { + "output_dim": int(self.action_space.n), + "num_discrete_choices": int(self.action_space.n), + } + ) elif isinstance(self.action_space, spaces.MultiDiscrete): assert self.action_space.nvec.ndim == 1, ( "Only one-dimensional MultiDiscrete action spaces are supported, " f"but found MultiDiscrete({(self.action_space.nvec).tolist()})." ) - self.actor_kwargs.update({ - "output_dim": int(np.sum(self.action_space.nvec)), - "num_discrete_choices": self.action_space.nvec, # type: ignore[dict-item] - }) + self.actor_kwargs.update( + { + "output_dim": int(np.sum(self.action_space.nvec)), + "num_discrete_choices": self.action_space.nvec, # type: ignore[dict-item] + } + ) elif isinstance(self.action_space, spaces.MultiBinary): assert isinstance(self.action_space.n, int), ( f"Multi-dimensional MultiBinary({self.action_space.n}) action space is not supported. " "You can flatten it instead." ) # Handle binary action spaces as discrete action spaces with two choices. - self.actor_kwargs.update({ - "output_dim": 2 * self.action_space.n, - "num_discrete_choices": 2 * np.ones(self.action_space.n, dtype=int), - }) + self.actor_kwargs.update( + { + "output_dim": 2 * self.action_space.n, + "num_discrete_choices": 2 * np.ones(self.action_space.n, dtype=int), + } + ) else: raise NotImplementedError(f"{self.action_space}") if self.featurizer_class is not None: - self.featurizer = self.featurizer_class( - **self.featurizer_kwargs - ) + self.featurizer = self.featurizer_class(**self.featurizer_kwargs) else: self.featurizer = Identity() - + optimizer = optax.inject_hyperparams(self.optimizer_class)( learning_rate=lr_schedule, **self.optimizer_kwargs ) @@ -286,9 +305,7 @@ def build(self, key: jax.Array, lr_schedule: Union[optax.Schedule, float], max_g obs = self.featurizer.apply(self.featurizer_state.params, obs) - self.actor = self.actor_class( - **self.actor_kwargs - ) + self.actor = self.actor_class(**self.actor_kwargs) self.actor_state = TrainState.create( apply_fn=self.actor.apply, @@ -299,9 +316,7 @@ def build(self, key: jax.Array, lr_schedule: Union[optax.Schedule, float], max_g ), ) - self.critic = self.critic_class( - **self.critic_kwargs - ) + self.critic = self.critic_class(**self.critic_kwargs) self.critic_state = TrainState.create( apply_fn=self.critic.apply, @@ -320,13 +335,19 @@ def build(self, key: jax.Array, lr_schedule: Union[optax.Schedule, float], max_g def reset_noise(self): self.key, self.noise_key = jr.split(self.key) - def forward(self, key: jax.Array, obs: np.ndarray, deterministic: bool = False) -> np.ndarray: + def forward( + self, key: jax.Array, obs: np.ndarray, deterministic: bool = False + ) -> np.ndarray: return self._predict(key, obs, deterministic=deterministic) def predict_all(self, key: jax.Array, observation: np.ndarray): - return self._predict_all(key, self.featurizer_state, self.actor_state, self.critic_state, observation) + return self._predict_all( + key, self.featurizer_state, self.actor_state, self.critic_state, observation + ) - def _predict(self, key: jax.Array, obs: np.ndarray, deterministic: bool = False) -> np.ndarray: + def _predict( + self, key: jax.Array, obs: np.ndarray, deterministic: bool = False + ) -> np.ndarray: if deterministic: return self.select_action(self.featurizer_state, self.actor_state, obs) else: @@ -347,7 +368,7 @@ def _predict_all(key, featurizer_state, actor_state, critic_state, observations) def select_action(featurizer_state, actor_state, obs): feats = featurizer_state.apply_fn(featurizer_state.params, obs) return actor_state.apply_fn(actor_state.params, feats).mode() - + @staticmethod @jit def sample_action(key, featurizer_state, actor_state, obs): @@ -356,8 +377,8 @@ def sample_action(key, featurizer_state, actor_state, obs): action = dist.sample(seed=key) return action -class BPOPolicy(PPOPolicy): +class BPOPolicy(PPOPolicy): def __init__( self, observation_space: spaces.Space, @@ -385,7 +406,7 @@ def __init__( actor_class=actor_class, actor_kwargs=actor_kwargs, critic_class=critic_class, - critic_kwargs=critic_kwargs + critic_kwargs=critic_kwargs, ) if q_vf_kwargs is None: @@ -400,9 +421,21 @@ def __init__( self.q_vf_kwargs = q_vf_kwargs - def build(self, key: jax.Array, lr_schedule: Union[optax.Schedule, float], max_grad_norm: float): + def build( + self, + key: jax.Array, + lr_schedule: Union[optax.Schedule, float], + max_grad_norm: float, + ): - key, feat_key, actor_key, critic_key, q_vf_key, q_hat_vf_key, = jax.random.split(key, 6) + ( + key, + feat_key, + actor_key, + critic_key, + q_vf_key, + q_hat_vf_key, + ) = jax.random.split(key, 6) key, self.key = jax.random.split(key, 2) self.reset_noise() @@ -414,43 +447,49 @@ def build(self, key: jax.Array, lr_schedule: Union[optax.Schedule, float], max_g act = act.reshape(-1, 1) if isinstance(self.action_space, spaces.Box): - self.actor_kwargs.update({ - "output_dim": int(np.prod(self.action_space.shape)), - }) + self.actor_kwargs.update( + { + "output_dim": int(np.prod(self.action_space.shape)), + } + ) elif isinstance(self.action_space, spaces.Discrete): - self.actor_kwargs.update({ - "output_dim": int(self.action_space.n), - "num_discrete_choices": int(self.action_space.n), - }) + self.actor_kwargs.update( + { + "output_dim": int(self.action_space.n), + "num_discrete_choices": int(self.action_space.n), + } + ) elif isinstance(self.action_space, spaces.MultiDiscrete): assert self.action_space.nvec.ndim == 1, ( "Only one-dimensional MultiDiscrete action spaces are supported, " f"but found MultiDiscrete({(self.action_space.nvec).tolist()})." ) - self.actor_kwargs.update({ - "output_dim": int(np.sum(self.action_space.nvec)), - "num_discrete_choices": self.action_space.nvec, # type: ignore[dict-item] - }) + self.actor_kwargs.update( + { + "output_dim": int(np.sum(self.action_space.nvec)), + "num_discrete_choices": self.action_space.nvec, # type: ignore[dict-item] + } + ) elif isinstance(self.action_space, spaces.MultiBinary): assert isinstance(self.action_space.n, int), ( f"Multi-dimensional MultiBinary({self.action_space.n}) action space is not supported. " "You can flatten it instead." ) # Handle binary action spaces as discrete action spaces with two choices. - self.actor_kwargs.update({ - "output_dim": 2 * self.action_space.n, - "num_discrete_choices": 2 * np.ones(self.action_space.n, dtype=int), - }) + self.actor_kwargs.update( + { + "output_dim": 2 * self.action_space.n, + "num_discrete_choices": 2 * np.ones(self.action_space.n, dtype=int), + } + ) else: raise NotImplementedError(f"{self.action_space}") if self.featurizer_class is not None: - self.featurizer = self.featurizer_class( - **self.featurizer_kwargs - ) + self.featurizer = self.featurizer_class(**self.featurizer_kwargs) else: self.featurizer = Identity() - + optimizer = optax.inject_hyperparams(self.optimizer_class)( learning_rate=lr_schedule, **self.optimizer_kwargs ) @@ -468,9 +507,7 @@ def build(self, key: jax.Array, lr_schedule: Union[optax.Schedule, float], max_g obs = self.featurizer.apply(self.featurizer_state.params, obs) - self.actor = self.actor_class( - **self.actor_kwargs - ) + self.actor = self.actor_class(**self.actor_kwargs) self.actor_state = TrainState.create( apply_fn=self.actor.apply, @@ -481,9 +518,7 @@ def build(self, key: jax.Array, lr_schedule: Union[optax.Schedule, float], max_g ), ) - self.critic = self.critic_class( - **self.critic_kwargs - ) + self.critic = self.critic_class(**self.critic_kwargs) self.critic_state = TrainState.create( apply_fn=self.critic.apply, @@ -497,9 +532,7 @@ def build(self, key: jax.Array, lr_schedule: Union[optax.Schedule, float], max_g self.actor.apply = jax.jit(self.actor.apply) self.critic.apply = jax.jit(self.critic.apply) - self.mu = self.actor_class( - **self.actor_kwargs - ) + self.mu = self.actor_class(**self.actor_kwargs) self.mu_state = TrainState.create( apply_fn=self.mu.apply, @@ -521,10 +554,7 @@ def build(self, key: jax.Array, lr_schedule: Union[optax.Schedule, float], max_g self.mu.apply = jax.jit(self.mu.apply) - self.q_vf = self.q_vf_class( - output_dim=q_vf_output_dim, - **self.q_vf_kwargs - ) + self.q_vf = self.q_vf_class(output_dim=q_vf_output_dim, **self.q_vf_kwargs) self.q_vf_state = ValueTrainState.create( apply_fn=self.q_vf.apply, @@ -533,13 +563,11 @@ def build(self, key: jax.Array, lr_schedule: Union[optax.Schedule, float], max_g optax.clip_by_global_norm(max_grad_norm), optimizer, ), - target_params=self.q_vf.init(q_vf_key, *q_vf_input) + target_params=self.q_vf.init(q_vf_key, *q_vf_input), ) self.q_hat_vf = self.q_vf_class( - softplus_final=True, - output_dim=q_vf_output_dim, - **self.q_vf_kwargs + softplus_final=True, output_dim=q_vf_output_dim, **self.q_vf_kwargs ) self.q_hat_vf_state = ValueTrainState.create( @@ -549,7 +577,7 @@ def build(self, key: jax.Array, lr_schedule: Union[optax.Schedule, float], max_g optax.clip_by_global_norm(max_grad_norm), optimizer, ), - target_params=self.q_hat_vf.init(q_hat_vf_key, *q_vf_input) + target_params=self.q_hat_vf.init(q_hat_vf_key, *q_vf_input), ) self.q_vf.apply = jax.jit(self.q_vf.apply) @@ -558,11 +586,20 @@ def build(self, key: jax.Array, lr_schedule: Union[optax.Schedule, float], max_g return key def predict_all(self, key: jax.Array, observation: np.ndarray): - return self._predict_all(key, self.featurizer_state, self.actor_state, self.mu_state, self.critic_state, observation) + return self._predict_all( + key, + self.featurizer_state, + self.actor_state, + self.mu_state, + self.critic_state, + observation, + ) @staticmethod @jit - def _predict_all(key, featurizer_state, actor_state, mu_state, vf_state, observations): + def _predict_all( + key, featurizer_state, actor_state, mu_state, vf_state, observations + ): features = featurizer_state.apply_fn(featurizer_state.params, observations) dist1 = mu_state.apply_fn(mu_state.params, features) dist2 = actor_state.apply_fn(actor_state.params, features) @@ -571,8 +608,3 @@ def _predict_all(key, featurizer_state, actor_state, mu_state, vf_state, observa actor_log_probs = dist2.log_prob(actions) values = vf_state.apply_fn(vf_state.params, features).flatten() return actions, mu_log_probs, actor_log_probs, values - - - - - \ No newline at end of file diff --git a/common/running_mean_std.py b/bpo/common/running_mean_std.py similarity index 88% rename from common/running_mean_std.py rename to bpo/common/running_mean_std.py index 93f235e..9623e15 100644 --- a/common/running_mean_std.py +++ b/bpo/common/running_mean_std.py @@ -1,8 +1,7 @@ - import numpy as np -class RunningMeanStd: +class RunningMeanStd: def __init__(self, eps=1e-8, shape=()): self.mean = np.zeros(shape, np.float64) self.var = np.zeros(shape, np.float64) @@ -29,7 +28,11 @@ def update_from_moments(self, batch_mean, batch_var, batch_count): new_mean = self.mean + delta * batch_count / total m_a = self.var * self.count m_b = batch_var * batch_count - m_2 = m_a + m_b + np.square(delta) * self.count * batch_count / (self.count + batch_count) + m_2 = ( + m_a + + m_b + + np.square(delta) * self.count * batch_count / (self.count + batch_count) + ) new_var = m_2 / (self.count + batch_count) new_count = batch_count + self.count @@ -37,4 +40,3 @@ def update_from_moments(self, batch_mean, batch_var, batch_count): self.mean = new_mean self.var = new_var self.count = new_count - diff --git a/common/symlog.py b/bpo/common/symlog.py similarity index 83% rename from common/symlog.py rename to bpo/common/symlog.py index 4c18b65..fea8740 100644 --- a/common/symlog.py +++ b/bpo/common/symlog.py @@ -2,16 +2,20 @@ import jax.numpy as jnp from jax import jit + def symlog(x): return np.sign(x) * np.log(1 + np.abs(x)) + def symexp(x): return np.sign(x) * (np.exp(np.abs(x)) - 1) + @jit def jsymlog(x): return jnp.sign(x) * jnp.log(1 + jnp.abs(x)) + @jit def jsymexp(x): - return jnp.sign(x) * (jnp.exp(jnp.abs(x)) - 1) \ No newline at end of file + return jnp.sign(x) * (jnp.exp(jnp.abs(x)) - 1) diff --git a/common/wrappers.py b/bpo/common/wrappers.py similarity index 94% rename from common/wrappers.py rename to bpo/common/wrappers.py index 7379c88..a5d394f 100644 --- a/common/wrappers.py +++ b/bpo/common/wrappers.py @@ -1,7 +1,7 @@ from __future__ import annotations import gymnasium as gym from gymnasium import spaces -from common.running_mean_std import RunningMeanStd +from bpo.common.running_mean_std import RunningMeanStd import numpy as np @@ -48,6 +48,7 @@ def is_wrapped(env: gym.Env, wrapper_class: gym.Wrapper) -> bool: return False + def get_wrapped(env: gym.Env, wrapper_class: gym.Wrapper) -> gym.Env: r""" Return the first wrapper instance of type ``wrapper_class`` found in ``env``'s @@ -85,6 +86,7 @@ def get_wrapped(env: gym.Env, wrapper_class: gym.Wrapper) -> gym.Env: return None + class ObsWrapper(gym.Wrapper): """ Base class for wrappers that transform observations while preserving constraint access. @@ -144,6 +146,7 @@ def step(self, action): obs, rew, term, trunc, info = self.env.step(action) return self._get_obs(obs), rew, term, trunc, info + class TimeLimit(gym.Wrapper): """ Episode time-limit wrapper compatible with constraint persistence. @@ -199,6 +202,7 @@ def reset(self, **kwargs): self._elapsed_steps = 0 return self.env.reset(**kwargs) + class RewardMonitor(gym.Wrapper): """ Monitor that injects reward/length metrics into ``info``. @@ -256,7 +260,7 @@ def step(self, action): if terminated or truncated: info["metrics"]["episode"] = self._episode_metrics() return observation, reward, terminated, truncated, info - + def _episode_metrics(self): """ Compute episode-level reward metrics. @@ -269,6 +273,7 @@ def _episode_metrics(self): """ return {"ep_reward": self.total_reward, "ep_length": self.total_steps} + class NormWrapper(gym.Wrapper): """ Normalize observations and/or rewards for a *single* (non-vectorized) environment. @@ -305,23 +310,23 @@ class NormWrapper(gym.Wrapper): """ def __init__( - self, - env: gym.Env, + self, + env: gym.Env, norm_obs: bool = True, norm_rew: bool = True, training: bool = True, clip_obs: float = 10.0, clip_rew: float = 10.0, gamma: float = 0.99, - eps: float = 1e-8 + eps: float = 1e-8, ): - assert not isinstance( - env, VecEnvWrapperBase - ), "NormWrapper does not expect a vectorized environment (DummyVecWrapper / VecWrapper). Please use VecNormWrapper instead" + assert not isinstance(env, VecEnvWrapperBase), ( + "NormWrapper does not expect a vectorized environment (DummyVecWrapper / VecWrapper). Please use VecNormWrapper instead" + ) - assert norm_obs and isinstance( - env.observation_space, spaces.Box - ), "NormWrapper only supports Box observation spaces when norm_obs=True." + assert norm_obs and isinstance(env.observation_space, spaces.Box), ( + "NormWrapper only supports Box observation spaces when norm_obs=True." + ) super().__init__(env) @@ -351,7 +356,7 @@ def _normalize_obs(self, obs: np.ndarray) -> np.ndarray: return np.clip( (obs - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.eps), -self.clip_obs, - self.clip_obs + self.clip_obs, ) def _normalize_rew(self, rew: float) -> float: @@ -422,6 +427,7 @@ def step(self, action): return obs, rew, term, trunc, info + class OneHotObsWrapper(ObsWrapper): """ One-hot encode :class:`gymnasium.spaces.Discrete` observations. @@ -457,7 +463,6 @@ def __init__(self, env: gym.Env): shape=(n,), dtype=np.float32, ) - elif isinstance(self._orig_obs_space, spaces.Dict): self._mode = "dict" @@ -532,10 +537,11 @@ def _get_obs(self, obs: Union[int, Dict[str, Any], np.ndarray]) -> np.ndarray: new_obs[key] = value return new_obs - else: + else: # pass return obs + class FlattenDictObsWrapper(ObsWrapper): """ Flatten a :class:`gymnasium.spaces.Dict` observation into a 1D :class:`~gymnasium.spaces.Box`. @@ -633,8 +639,8 @@ def _get_obs(self, obs: Dict[str, Any]) -> np.ndarray: this implementation currently enforces Box-only subspaces at runtime. """ assert isinstance(obs, dict), ( - f"Expected dict observation for Dict space, got {type(obs)}" - ) + f"Expected dict observation for Dict space, got {type(obs)}" + ) parts = [] for key in sorted(self._orig_obs_space.spaces.keys()): @@ -652,6 +658,7 @@ def _get_obs(self, obs: Dict[str, Any]) -> np.ndarray: return np.concatenate(parts, axis=0).astype(np.float32) + class VecEnvWrapperBase(gym.Wrapper): """ Base class for simple Python-list vector environment wrappers. @@ -681,11 +688,11 @@ def __init__(self, env: gym.Env): super().__init__(env) def reset_done( - self, + self, dones: Union[List[bool], np.ndarray], - *, - seed: int | None = None, - options: Dict[str, Any] | None = None + *, + seed: int | None = None, + options: Dict[str, Any] | None = None, ): """ Reset only the environments indicated by ``dones``. @@ -709,6 +716,7 @@ def reset_done( """ raise NotImplementedError + class DummyVecWrapper(VecEnvWrapperBase): """ Wrap a single environment with a vector-environment API (``n_envs=1``). @@ -744,11 +752,11 @@ def reset(self, *, seed: int | None = None, options: Dict[str, Any] | None = Non return [obs], [info] def reset_done( - self, + self, dones: Union[List[bool], np.ndarray], - *, - seed: int | None = None, - options: Dict[str, Any] | None = None + *, + seed: int | None = None, + options: Dict[str, Any] | None = None, ): """ Conditionally reset the single environment. @@ -782,6 +790,7 @@ def step(self, action): obs, rew, term, trunc, info = self.env.step(action) return [obs], [rew], [term], [trunc], [info] + class VecWrapper(VecEnvWrapperBase): """ Wrap a list of environments with a simple vector-environment API. @@ -799,7 +808,7 @@ class VecWrapper(VecEnvWrapperBase): def __init__(self, envs: List[gym.Env]): assert len(envs) > 0, "VecWrapper requires at least one environment" - super().__init__(envs[0]) # maintain API compatibility + super().__init__(envs[0]) # maintain API compatibility self.envs: List[gym.Env] = envs self.n_envs = len(envs) @@ -823,11 +832,11 @@ def reset(self, *, seed: int | None = None, options: Dict[str, Any] | None = Non return obs_list, info_list def reset_done( - self, + self, dones: Union[List[bool], np.ndarray], - *, - seed: int | None = None, - options: Dict[str, Any] | None = None + *, + seed: int | None = None, + options: Dict[str, Any] | None = None, ): """ Reset only environments whose done flag is True. @@ -881,6 +890,7 @@ def step(self, actions): return obs_list, rew_list, term_list, trunc_list, info_list + class VecNormWrapper(VecEnvWrapperBase): """ Normalize observations and/or rewards for a vectorized environment. @@ -911,23 +921,23 @@ class VecNormWrapper(VecEnvWrapperBase): """ def __init__( - self, - env: Union[gym.Env, List[gym.Env]], + self, + env: Union[gym.Env, List[gym.Env]], norm_obs: bool = True, norm_rew: bool = True, training: bool = True, clip_obs: float = 10.0, clip_rew: float = 10.0, gamma: float = 0.99, - eps: float = 1e-8 + eps: float = 1e-8, ): - assert isinstance( - env, VecEnvWrapperBase - ), "VecNormWrapper expects a vectorized environment (DummyVecWrapper / VecWrapper)." + assert isinstance(env, VecEnvWrapperBase), ( + "VecNormWrapper expects a vectorized environment (DummyVecWrapper / VecWrapper)." + ) - assert norm_obs and isinstance( - env.observation_space, spaces.Box - ), "VecNormWrapper only supports Box observation spaces when norm_obs=True." + assert norm_obs and isinstance(env.observation_space, spaces.Box), ( + "VecNormWrapper only supports Box observation spaces when norm_obs=True." + ) super().__init__(env) @@ -999,11 +1009,11 @@ def reset(self, *, seed: int | None = None, options: Dict[str, Any] | None = Non return obs_list, info_list def reset_done( - self, + self, dones: Union[List[bool], np.ndarray], - *, - seed: int | None = None, - options: Dict[str, Any] | None = None + *, + seed: int | None = None, + options: Dict[str, Any] | None = None, ): """ Reset only environments indicated by ``dones`` and normalize those observations. @@ -1018,14 +1028,16 @@ def reset_done( :meth:`VecEnvWrapperBase.reset_done`, with reset observations optionally normalized. """ - reset_obs, reset_infos = self.env.reset_done( - dones, seed=seed, options=options - ) + reset_obs, reset_infos = self.env.reset_done(dones, seed=seed, options=options) - obs_arr = np.asarray( - [o for o in reset_obs if o is not None], - dtype=np.float32, - ) if any(o is not None for o in reset_obs) else None + obs_arr = ( + np.asarray( + [o for o in reset_obs if o is not None], + dtype=np.float32, + ) + if any(o is not None for o in reset_obs) + else None + ) if self.norm_obs and self.training and obs_arr is not None: self.obs_rms.update(obs_arr) @@ -1040,7 +1052,9 @@ def reset_done( for i, done in enumerate(dones): if done and reset_obs[i] is not None: o = np.asarray(reset_obs[i], dtype=np.float32) - norm = (o - self.obs_rms.mean) / np.sqrt(self.obs_rms.var + self.eps) + norm = (o - self.obs_rms.mean) / np.sqrt( + self.obs_rms.var + self.eps + ) norm = np.clip(norm, -self.clip_obs, self.clip_obs) norm_reset_obs[i] = norm reset_obs = norm_reset_obs @@ -1076,4 +1090,4 @@ def step(self, actions): if self.norm_obs: obs_list = self._normalize_obs(obs_list) - return obs_list, rew_list, term_list, trunc_list, infos \ No newline at end of file + return obs_list, rew_list, term_list, trunc_list, infos diff --git a/configs.yaml b/bpo/configs.yaml similarity index 100% rename from configs.yaml rename to bpo/configs.yaml diff --git a/bpo/envs/__init__.py b/bpo/envs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/envs/cartpole.py b/bpo/envs/cartpole.py similarity index 84% rename from envs/cartpole.py rename to bpo/envs/cartpole.py index 846e10c..3f40c6d 100644 --- a/envs/cartpole.py +++ b/bpo/envs/cartpole.py @@ -10,9 +10,9 @@ import numpy as np from typing import Dict, Any -class Cartpole(gym.Env): - def __init__(self, continuous: bool =False): +class Cartpole(gym.Env): + def __init__(self, continuous: bool = False): self._gravity = 9.8 self._masscart = 1.0 @@ -32,18 +32,20 @@ def __init__(self, continuous: bool =False): high = np.array( [ - self._x_threshold*2, - self._x_vel_threshold*2, - self._theta_threshold_radians*2, - self._theta_vel_threshold*2, + self._x_threshold * 2, + self._x_vel_threshold * 2, + self._theta_threshold_radians * 2, + self._theta_vel_threshold * 2, ], - dtype=np.float32 + dtype=np.float32, ) self.observation_space = spaces.Box(low=-high, high=high, dtype=np.float32) if continuous: - self.action_space = spaces.Box(low=-1.0, high=1.0, shape=(1,), dtype=np.float32) + self.action_space = spaces.Box( + low=-1.0, high=1.0, shape=(1,), dtype=np.float32 + ) else: self.action_space = gym.spaces.Discrete(2) @@ -61,9 +63,9 @@ def reset(self, *, seed: int | None = None, options: Dict[str, Any] | None = Non self.np_random = np.random.default_rng(seed) self._state = self.np_random.uniform( - low=np.array([-0.05, -0.05, -0.05, -0.05]), - high=np.array([0.05, 0.05, 0.05, 0.05]), - size=(4,) + low=np.array([-0.05, -0.05, -0.05, -0.05]), + high=np.array([0.05, 0.05, 0.05, 0.05]), + size=(4,), ) return self._obs(), {} @@ -74,8 +76,10 @@ def step(self, action: Any): x, x_dot, theta, theta_dot = self._state - stable = np.abs(theta) <= self._theta_threshold_radians \ + stable = ( + np.abs(theta) <= self._theta_threshold_radians and np.abs(x) <= self._x_threshold + ) return self._obs(), 1.0, bool(not stable), False, {} @@ -84,7 +88,6 @@ def _step(self, state: Any, action: Any) -> np.ndarray: class DiscreteCartpole(Cartpole): - def __init__(self, **kwargs): super().__init__(continuous=False) @@ -98,10 +101,10 @@ def _step(self, state: Any, action: Any) -> np.ndarray: # For the interested reader: # https://coneural.org/florian/papers/05_cart_pole.pdf temp = ( - force + self._polemass_length * theta_dot ** 2 * sintheta + force + self._polemass_length * theta_dot**2 * sintheta ) / self._total_mass thetaacc = (self._gravity * sintheta - costheta * temp) / ( - self._length * (4.0 / 3.0 - self._masspole * costheta ** 2 / self._total_mass) + self._length * (4.0 / 3.0 - self._masspole * costheta**2 / self._total_mass) ) xacc = temp - self._polemass_length * thetaacc * costheta / self._total_mass @@ -120,7 +123,6 @@ def _step(self, state: Any, action: Any) -> np.ndarray: class ContinuousCartpole(Cartpole): - def __init__(self, **kwargs): super().__init__(continuous=True) @@ -134,10 +136,10 @@ def _step(self, state, action) -> np.ndarray: # For the interested reader: # https://coneural.org/florian/papers/05_cart_pole.pdf temp = ( - force + self._polemass_length * theta_dot ** 2 * sintheta + force + self._polemass_length * theta_dot**2 * sintheta ) / self._total_mass thetaacc = (self._gravity * sintheta - costheta * temp) / ( - self._length * (4.0 / 3.0 - self._masspole * costheta ** 2 / self._total_mass) + self._length * (4.0 / 3.0 - self._masspole * costheta**2 / self._total_mass) ) xacc = temp - self._polemass_length * thetaacc * costheta / self._total_mass @@ -153,5 +155,3 @@ def _step(self, state, action) -> np.ndarray: theta = theta + self._tau * theta_dot return np.array([x, x_dot, theta, theta_dot]) - - \ No newline at end of file diff --git a/envs/mujoco.py b/bpo/envs/mujoco.py similarity index 86% rename from envs/mujoco.py rename to bpo/envs/mujoco.py index f41aaf4..16dacf6 100644 --- a/envs/mujoco.py +++ b/bpo/envs/mujoco.py @@ -1,11 +1,10 @@ - import math import gymnasium as gym import numpy as np from typing import Dict, Any -class MuJoCoEnv(gym.Env): +class MuJoCoEnv(gym.Env): def __init__(self, env: gym.Env): self._env = env self.observation_space = self._env.observation_space @@ -21,33 +20,26 @@ def step(self, action: Any): assert self.action_space.contains(action), err_msg return self._env.step(action) -class Ant(MuJoCoEnv): +class Ant(MuJoCoEnv): def __init__(self, **kwargs): - env = gym.make('Ant-v5') + env = gym.make("Ant-v5") super().__init__(env) -class HalfCheetah(MuJoCoEnv): +class HalfCheetah(MuJoCoEnv): def __init__(self, **kwargs): - env = gym.make('HalfCheetah-v5') + env = gym.make("HalfCheetah-v5") super().__init__(env) -class Hopper(MuJoCoEnv): +class Hopper(MuJoCoEnv): def __init__(self, **kwargs): - env = gym.make('Hopper-v5') + env = gym.make("Hopper-v5") super().__init__(env) -class Walker2D(MuJoCoEnv): +class Walker2D(MuJoCoEnv): def __init__(self, **kwargs): - env = gym.make('Walker2d-v5') + env = gym.make("Walker2d-v5") super().__init__(env) - - - - - - - diff --git a/ppo.py b/bpo/ppo.py similarity index 81% rename from ppo.py rename to bpo/ppo.py index cca0ca9..52b048f 100644 --- a/ppo.py +++ b/bpo/ppo.py @@ -11,14 +11,14 @@ import gymnasium as gym from gymnasium import spaces from typing import Any, Optional, TypeVar, Union, Callable -from common.base_class import BaseJaxPolicy -from common.on_policy_algorithm import OnPolicyAlgorithm -from common.policies import PPOPolicy -from common.metrics import StatsLogger, Stats +from bpo.common.base_class import BaseJaxPolicy +from bpo.common.on_policy_algorithm import OnPolicyAlgorithm +from bpo.common.policies import PPOPolicy +from bpo.common.metrics import StatsLogger, Stats from tqdm.auto import tqdm -class PPO(OnPolicyAlgorithm): +class PPO(OnPolicyAlgorithm): def __init__( self, env: gym.Env, @@ -28,7 +28,7 @@ def __init__( monitor: bool = True, device: str = "auto", verbose: int = 0, - eval_env: Optional[gym.Env] = None, + eval_env: Optional[gym.Env] = None, learning_rate: Union[float, optax.Schedule] = 3e-4, n_steps: int = 2048, batch_size: int = 64, @@ -48,14 +48,14 @@ def __init__( ): super().__init__( - env, + env, tensorboard_logdir=tensorboard_logdir, seed=seed, monitor=monitor, device=device, verbose=verbose, eval_env=eval_env, - use_tqdm_rollout=True, # Turn on tqdm progress bar for rollout + use_tqdm_rollout=True, # Turn on tqdm progress bar for rollout learning_rate=learning_rate, n_steps=n_steps, gamma=gamma, @@ -73,7 +73,9 @@ def __init__( if isinstance(clip_range, float): self.clip_range_schedule = optax.schedules.constant_schedule(clip_range) else: - assert callable(clip_range), f"clip_range for class PPO must be float or optax.Schedule not {clip_range}" + assert callable(clip_range), ( + f"clip_range for class PPO must be float or optax.Schedule not {clip_range}" + ) self.clip_range_schedule = clip_range self.batch_size = batch_size @@ -121,14 +123,14 @@ def actor_critic_loss(featurizer_params, actor_params, critic_params): # Critic loss critic_values = critic_state.apply_fn(critic_params, features).flatten() - value_loss = vf_coef * ((returns - critic_values)**2).mean() + value_loss = vf_coef * ((returns - critic_values) ** 2).mean() total_loss = total_policy_loss + value_loss return total_loss, (total_policy_loss, value_loss) - (loss, (pg_loss, vf_loss)), grads = jax.value_and_grad(actor_critic_loss, argnums=(0, 1, 2), has_aux=True)( - featurizer_state.params, actor_state.params, critic_state.params - ) + (loss, (pg_loss, vf_loss)), grads = jax.value_and_grad( + actor_critic_loss, argnums=(0, 1, 2), has_aux=True + )(featurizer_state.params, actor_state.params, critic_state.params) featurizer_state = featurizer_state.apply_gradients(grads=grads[0]) actor_state = actor_state.apply_gradients(grads=grads[1]) @@ -136,12 +138,8 @@ def actor_critic_loss(featurizer_params, actor_params, critic_params): return (featurizer_state, actor_state, critic_state), (pg_loss, vf_loss) - def optimize( - self, - step: int, - logger: Optional[StatsLogger] = None, - tqdm_position: int = 1 + self, step: int, logger: Optional[StatsLogger] = None, tqdm_position: int = 1 ): if self.verbose > 0: @@ -159,26 +157,40 @@ def optimize( self.key, key1, _, _, _ = jr.split(self.key, 5) with tqdm( - total=self.n_epochs*self.n_steps//(self.batch_size//self.n_envs), + total=self.n_epochs * self.n_steps // (self.batch_size // self.n_envs), desc="optimize", position=tqdm_position, leave=False, dynamic_ncols=True, colour="cyan", ) as pbar: - for _ in range(self.n_epochs): key1, subkey = jr.split(key1) - for rollout_data in self.rollout_buffer.get(subkey, self.batch_size//self.n_envs): - - observations, actions, rewards, values, returns, advantages, old_log_probs = rollout_data + for rollout_data in self.rollout_buffer.get( + subkey, self.batch_size // self.n_envs + ): + ( + observations, + actions, + rewards, + values, + returns, + advantages, + old_log_probs, + ) = rollout_data if isinstance(self.action_space, spaces.Discrete): # Convert discrete action from float to int actions = actions.flatten().astype(np.int32) - (self.policy.featurizer_state, self.policy.actor_state, self.policy.critic_state), (pg_loss, vf_loss) = \ - self._one_update( + ( + ( + self.policy.featurizer_state, + self.policy.actor_state, + self.policy.critic_state, + ), + (pg_loss, vf_loss), + ) = self._one_update( featurizer_state=self.policy.featurizer_state, actor_state=self.policy.actor_state, critic_state=self.policy.critic_state, @@ -198,7 +210,7 @@ def optimize( if self.verbose > 0: pg_loss_stats.update(float(pg_loss)) vf_loss_stats.update(float(vf_loss)) - + if self.verbose > 1: rewards_stats.update(np.array(rewards)) values_stats.update(np.array(values)) @@ -237,8 +249,7 @@ def optimize( if logger: logger.add(logs) - + @property def train_ratio(self): return self.n_steps * self.n_envs - diff --git a/bpo/run.py b/bpo/run.py new file mode 100644 index 0000000..1402963 --- /dev/null +++ b/bpo/run.py @@ -0,0 +1,289 @@ +import warnings + +warnings.filterwarnings("ignore", category=DeprecationWarning) +warnings.filterwarnings("ignore", category=UserWarning) + +import os + +from bpo.ppo import PPO +from bpo.bpo import BPO_PPO +from bpo.common.configs import Config, Flags +from bpo.common.wrappers import ( + TimeLimit, + RewardMonitor, + VecNormWrapper, + NormWrapper, + VecWrapper, +) +from bpo.common.policies import PPOPolicy, BPOPolicy +from bpo.common.layers import NatureCNN + +import optax +from optax.schedules import linear_schedule +import flax.linen as nn +from flax.linen import initializers + +import importlib +import importlib.resources as resources +import math +import ruamel.yaml as yaml +from typing import Dict, Any + + +def make_env(env_id: str, max_episode_steps: int, **kwargs): + ctor = { + "cartpole": "bpo.envs.cartpole:DiscreteCartpole", + "cont_cartpole": "bpo.envs.cartpole:ContinuousCartpole", + "ant": "bpo.envs.mujoco:Ant", + "half_cheetah": "bpo.envs.mujoco:HalfCheetah", + "hopper": "bpo.envs.mujoco:Hopper", + "walker_2d": "bpo.envs.mujoco:Walker2D", + }[env_id] + if isinstance(ctor, str): + module, cls = ctor.split(":") + module = importlib.import_module(module) + ctor = getattr(module, cls) + env = ctor(**kwargs) + env = TimeLimit(env, max_episode_steps) + env = RewardMonitor(env) + return env + + +def make_envs(n_envs: int, env_id: str, max_episode_steps: int, **kwargs): + ctor = { + "cartpole": "bpo.envs.cartpole:DiscreteCartpole", + "cont_cartpole": "bpo.envs.cartpole:ContinuousCartpole", + "ant": "bpo.envs.mujoco:Ant", + "half_cheetah": "bpo.envs.mujoco:HalfCheetah", + "hopper": "bpo.envs.mujoco:Hopper", + "walker_2d": "bpo.envs.mujoco:Walker2D", + }[env_id] + if isinstance(ctor, str): + module, cls = ctor.split(":") + module = importlib.import_module(module) + ctor = getattr(module, cls) + envs = [ + RewardMonitor(TimeLimit(ctor(**kwargs), max_episode_steps)) + for _ in range(n_envs) + ] + return envs + + +def parse_policy_kwargs(config: Config): + + try: + algo_kwargs = Config(config[config.algo]) + except: + raise NotImplementedError(f"algorithm: {config.algo}") + + if config.bpo: + policy_class = BPOPolicy + else: + policy_class = PPOPolicy + + total_timesteps = config.run.total_timesteps + n_envs = config.run.n_envs + n_steps = algo_kwargs.n_steps + + if hasattr(algo_kwargs, "clip_range"): + if hasattr(algo_kwargs, "clip_decay") and algo_kwargs.clip_decay == "linear": + clip_schedule = linear_schedule( + float(algo_kwargs.clip_range), 0.0, total_timesteps + ) + else: + clip_schedule = float(algo_kwargs.clip_range) + else: + clip_schedule = None + + if hasattr(algo_kwargs, "batch_size"): + batch_size = algo_kwargs.batch_size + else: + batch_size = 1 + + if hasattr(algo_kwargs, "n_epochs"): + n_epochs = algo_kwargs.n_epochs + else: + n_epochs = 1 + + decay_steps = ( + (n_steps * n_envs // batch_size) + * n_epochs + * math.ceil(total_timesteps / (n_steps * n_envs)) + ) + + if algo_kwargs.use_featurizer: + featurizer_class = NatureCNN + else: + featurizer_class = None + + featurizer_kwargs = dict( + grayscale_obs=algo_kwargs.featurizer_kwargs.grayscale_obs, + normalize_images=algo_kwargs.featurizer_kwargs.normalize_images, + activation_fn={ + "tanh": nn.tanh, + "relu": nn.relu, + }[algo_kwargs.featurizer_kwargs.activation_fn], + kernel_init={ + "lecun_normal": initializers.lecun_normal(), + "orthogonal": initializers.orthogonal(), + }[algo_kwargs.featurizer_kwargs.kernel_init], + ) + + optimizer_class = {"adam": optax.adam, "rms_prop": optax.rmsprop}[ + algo_kwargs.optimizer.opt + ] + + if algo_kwargs.optimizer.decay == "none": + lr_schedule = algo_kwargs.optimizer.learning_rate + elif algo_kwargs.optimizer.decay == "linear": + lr_schedule = linear_schedule( + algo_kwargs.optimizer.learning_rate, 0.0, decay_steps + ) + else: + raise NotImplementedError(algo_kwargs.optimizer.decay) + + max_grad_norm = algo_kwargs.optimizer.max_grad_norm + + optimizer_kwargs = dict(eps=algo_kwargs.optimizer.eps) + + actor_kwargs = dict( + n_units=algo_kwargs.actor_kwargs.n_units, + log_std_init=algo_kwargs.actor_kwargs.log_std_init, + activation_fn={ + "tanh": nn.tanh, + "relu": nn.relu, + }[algo_kwargs.actor_kwargs.activation_fn], + kernel_init={ + "lecun_normal": initializers.lecun_normal(), + "orthogonal": initializers.orthogonal(), + }[algo_kwargs.actor_kwargs.kernel_init], + ) + + critic_kwargs = dict( + n_units=algo_kwargs.critic_kwargs.n_units, + use_layer_norm=algo_kwargs.critic_kwargs.use_layer_norm, + activation_fn={ + "tanh": nn.tanh, + "relu": nn.relu, + }[algo_kwargs.critic_kwargs.activation_fn], + kernel_init={ + "lecun_normal": initializers.lecun_normal(), + "orthogonal": initializers.orthogonal(), + }[algo_kwargs.critic_kwargs.kernel_init], + ) + + q_vf_kwargs = dict( + n_units=algo_kwargs.q_vf_kwargs.n_units, + use_layer_norm=algo_kwargs.q_vf_kwargs.use_layer_norm, + use_zero_norm_final=algo_kwargs.q_vf_kwargs.use_zero_norm_final, + activation_fn={ + "tanh": nn.tanh, + "relu": nn.relu, + }[algo_kwargs.q_vf_kwargs.activation_fn], + kernel_init={ + "lecun_normal": initializers.lecun_normal(), + "orthogonal": initializers.orthogonal(), + }[algo_kwargs.q_vf_kwargs.kernel_init], + ) + + policy_kwargs = dict( + featurizer_class=featurizer_class, + featurizer_kwargs=featurizer_kwargs, + optimizer_class=optimizer_class, + optimizer_kwargs=optimizer_kwargs, + actor_kwargs=actor_kwargs, + critic_kwargs=critic_kwargs, + q_vf_kwargs=q_vf_kwargs, + ) + + algo_kwargs = {**algo_kwargs} + algo_kwargs.update( + { + "learning_rate": lr_schedule, + "clip_range": clip_schedule, + "max_grad_norm": max_grad_norm, + "policy_class": policy_class, + "policy_kwargs": policy_kwargs, + } + ) + + return algo_kwargs + + +def main(argv=None): + parsed, other = Flags(configs=["defaults"]).parse_known(argv) + # importlib.resources reads the packaged file as UTF-8. The custom Path + # class only splits on "/", so Path(__file__) misses configs.yaml on Windows. + configs = yaml.YAML(typ="safe").load( + resources.files("bpo").joinpath("configs.yaml").read_text(encoding="utf-8") + ) + config = Config(configs["defaults"]) + for name in parsed.configs: + config = config.update(configs[name]) + config = Flags(config).parse(other) + + print(config) + + try: + algo_kwargs = Config(config[config.algo]) + except: + raise NotImplementedError(f"algorithm: {config.algo}") + + if config.run.n_envs > 1: + envs = make_envs( + config.run.n_envs, config.env.env_id, config.env.max_episode_steps + ) + if config.env.normalize_obs or config.env.normalize_rew: + env = VecNormWrapper( + envs, + norm_obs=config.env.normalize_obs, + norm_rew=config.env.normalize_rew, + training=True, + gamma=algo_kwargs.gamma, + ) + else: + env = VecWrapper(envs) + else: + env = make_env(config.env.env_id, config.env.max_episode_steps) + if config.env.normalize_obs or config.env.normalize_rew: + env = NormWrapper( + env, + norm_obs=config.env.normalize_obs, + norm_rew=config.env.normalize_rew, + training=True, + gamma=algo_kwargs.gamma, + ) + + eval_env = make_env(config.env.env_id, config.env.max_episode_steps) + + base_kwargs = dict( + tensorboard_logdir=config.run.logdir, + seed=config.run.seed, + device=config.run.device, + verbose=config.run.verbose, + eval_env=eval_env, + ) + + algo_id = f"{'bpo_' if bool(config.bpo) else ''}{config.algo}" + + algo_class = { + "ppo": PPO, + "bpo_ppo": BPO_PPO, + }[algo_id] + + algo_kwargs = parse_policy_kwargs(config) + merged = {**base_kwargs, **algo_kwargs} + algo = algo_class(env, **merged) + + run_kwargs = dict( + num_eval_episodes=config.run.eval_episodes, + eval_freq=config.run.eval_every, + log_freq=config.run.log_every, + stats_window_size=config.run.stats_window_size, + ) + + algo.train(config.run.total_timesteps, **run_kwargs) + + +if __name__ == "__main__": + main() diff --git a/common/configs.py b/common/configs.py deleted file mode 100644 index 0807c35..0000000 --- a/common/configs.py +++ /dev/null @@ -1,511 +0,0 @@ -import re -import sys -import io -import json - -import contextlib -import glob -import os -import re -import shutil - - -class Path: - - filesystems = [] - - def __new__(cls, path): - path = str(path) - for impl, pred in cls.filesystems: - if pred(path): - obj = super().__new__(impl) - obj.__init__(path) - return obj - raise NotImplementedError(f'No filesystem supports: {path}') - - def __getnewargs__(self): - return (self._path,) - - def __init__(self, path): - assert isinstance(path, str) - path = re.sub(r'^\./*', '', path) # Remove leading dot or dot slashes. - path = re.sub(r'(?<=[^/])/$', '', path) # Remove single trailing slash. - path = path or '.' # Empty path is represented by a dot. - self._path = path - - def __truediv__(self, part): - sep = '' if self._path.endswith('/') else '/' - return type(self)(f'{self._path}{sep}{str(part)}') - - def __repr__(self): - return f'Path({str(self)})' - - def __fspath__(self): - return str(self) - - def __eq__(self, other): - return self._path == other._path - - def __lt__(self, other): - return self._path < other._path - - def __str__(self): - return self._path - - @property - def parent(self): - if '/' not in self._path: - return type(self)('.') - parent = self._path.rsplit('/', 1)[0] - parent = parent or ('/' if self._path.startswith('/') else '.') - return type(self)(parent) - - @property - def name(self): - if '/' not in self._path: - return self._path - return self._path.rsplit('/', 1)[1] - - @property - def stem(self): - return self.name.split('.', 1)[0] if '.' in self.name else self.name - - @property - def suffix(self): - return ('.' + self.name.split('.', 1)[1]) if '.' in self.name else '' - - def read(self, mode='r'): - assert mode in 'r rb'.split(), mode - with self.open(mode) as f: - return f.read() - - def write(self, content, mode='w'): - assert mode in 'w a wb ab'.split(), mode - with self.open(mode) as f: - f.write(content) - - @contextlib.contextmanager - def open(self, mode='r'): - raise NotImplementedError - - def absolute(self): - raise NotImplementedError - - def glob(self, pattern): - raise NotImplementedError - - def exists(self): - raise NotImplementedError - - def isfile(self): - raise NotImplementedError - - def isdir(self): - raise NotImplementedError - - def mkdirs(self): - raise NotImplementedError - - def remove(self): - raise NotImplementedError - - def rmtree(self): - raise NotImplementedError - - def copy(self, dest): - raise NotImplementedError - - def move(self, dest): - self.copy(dest) - self.remove() - - -class LocalPath(Path): - - def __init__(self, path): - super().__init__(os.path.expanduser(str(path))) - - @contextlib.contextmanager - def open(self, mode='r'): - with open(str(self), mode=mode) as f: - yield f - - def absolute(self): - return type(self)(os.path.absolute(str(self))) - - def glob(self, pattern): - for path in glob.glob(f'{str(self)}/{pattern}'): - yield type(self)(path) - - def exists(self): - return os.path.exists(str(self)) - - def isfile(self): - return os.path.isfile(str(self)) - - def isdir(self): - return os.path.isdir(str(self)) - - def mkdirs(self): - os.makedirs(str(self), exist_ok=True) - - def remove(self): - os.rmdir(str(self)) if self.isdir() else os.remove(str(self)) - - def rmtree(self): - shutil.rmtree(self) - - def copy(self, dest): - if self.isfile(): - shutil.copy(self, type(self)(dest)) - else: - shutil.copytree(self, type(self)(dest), dirs_exist_ok=True) - - def move(self, dest): - shutil.move(self, dest) - - -class GFilePath(Path): - - def __init__(self, path): - path = str(path) - if not (path.startswith('/') or '://' in path): - path = os.path.abspath(os.path.expanduser(path)) - super().__init__(path) - import tensorflow as tf - self._gfile = tf.io.gfile - - @contextlib.contextmanager - def open(self, mode='r'): - path = str(self) - if 'a' in mode and path.startswith('/cns/'): - path += '%r=3.2' - if mode.startswith('x') and self.exists(): - raise FileExistsError(path) - mode = mode.replace('x', 'w') - with self._gfile.GFile(path, mode) as f: - yield f - - def absolute(self): - return self - - def glob(self, pattern): - for path in self._gfile.glob(f'{str(self)}/{pattern}'): - yield type(self)(path) - - def exists(self): - return self._gfile.exists(str(self)) - - def isfile(self): - return self.exists() and not self.isdir() - - def isdir(self): - return self._gfile.isdir(str(self)) - - def mkdirs(self): - self._gfile.makedirs(str(self)) - - def remove(self): - self._gfile.remove(str(self)) - - def rmtree(self): - self._gfile.rmtree(str(self)) - - def copy(self, dest): - self._gfile.copy(str(self), str(dest), overwrite=True) - - def move(self, dest): - dest = Path(dest) - if dest.isdir(): - dest.rmtree() - self._gfile.rename(self, str(dest), overwrite=True) - - -Path.filesystems = [ - (GFilePath, lambda path: path.startswith('gs://')), - (GFilePath, lambda path: path.startswith('/cns/')), - (LocalPath, lambda path: True), -] - -class Flags: - - def __init__(self, *args, **kwargs): - self._config = Config(*args, **kwargs) - - def parse(self, argv=None, help_exists=True): - parsed, remaining = self.parse_known(argv) - for flag in remaining: - if flag.startswith('--'): - raise ValueError(f"Flag '{flag}' did not match any config keys.") - #assert not remaining or , remaining - return parsed - - def parse_known(self, argv=None, help_exists=False): - if argv is None: - argv = sys.argv[1:] - if '--help' in argv: - print('\nHelp:') - lines = str(self._config).split('\n')[2:] - print('\n'.join('--' + re.sub(r'[:,\[\]]', '', x) for x in lines)) - help_exists and sys.exit() - parsed = {} - remaining = [] - key = None - vals = None - for arg in argv: - if arg.startswith('--'): - if key: - self._submit_entry(key, vals, parsed, remaining) - if '=' in arg: - key, val = arg.split('=', 1) - vals = [val] - else: - key, vals = arg, [] - else: - if key: - vals.append(arg) - else: - remaining.append(arg) - self._submit_entry(key, vals, parsed, remaining) - parsed = self._config.update(parsed) - return parsed, remaining - - def _submit_entry(self, key, vals, parsed, remaining): - if not key and not vals: - return - if not key: - vals = ', '.join(f"'{x}'" for x in vals) - raise ValueError(f"Values {vals} were not preceded by any flag.") - name = key[len('--'):] - if '=' in name: - remaining.extend([key] + vals) - return - if self._config.IS_PATTERN.fullmatch(name): - pattern = re.compile(name) - keys = {k for k in self._config.flat if pattern.fullmatch(k)} - elif name in self._config: - keys = [name] - else: - keys = [] - if not keys: - remaining.extend([key] + vals) - return - if not vals: - raise ValueError(f"Flag '{key}' was not followed by any values.") - for key in keys: - parsed[key] = self._parse_flag_value(self._config[key], vals, key) - - def _parse_flag_value(self, default, value, key): - value = value if isinstance(value, (tuple, list)) else (value,) - if isinstance(default, (tuple, list)): - if len(value) == 1 and ',' in value[0]: - value = value[0].split(',') - return tuple(self._parse_flag_value(default[0], [x], key) for x in value) - assert len(value) == 1, value - value = str(value[0]) - if default is None: - return value - if isinstance(default, bool): - try: - return bool(['False', 'True'].index(value)) - except ValueError: - message = f"Expected bool but got '{value}' for key '{key}'." - raise TypeError(message) - if isinstance(default, int): - try: - value = float(value) # Allow scientific notation for integers. - assert float(int(value)) == value - except (TypeError, AssertionError): - message = f"Expected int but got float '{value}' for key '{key}'." - raise TypeError(message) - return int(value) - if isinstance(default, dict): - raise TypeError( - f"Key '{key}' refers to a whole dict. Please speicfy a subkey.") - return type(default)(value) - - -class Config(dict): - - SEP = '.' - IS_PATTERN = re.compile(r'.*[^A-Za-z0-9_.-].*') - - def __init__(self, *args, **kwargs): - mapping = dict(*args, **kwargs) - mapping = self._flatten(mapping) - mapping = self._ensure_keys(mapping) - mapping = self._ensure_values(mapping) - self._flat = mapping - self._nested = self._nest(mapping) - # Need to assign the values to the base class dictionary so that - # conversion to dict does not lose the content. - super().__init__(self._nested) - - @property - def flat(self): - return self._flat.copy() - - def save(self, filename): - filename = path.Path(filename) - if filename.suffix == '.json': - filename.write(json.dumps(dict(self))) - elif filename.suffix in ('.yml', '.yaml'): - import ruamel.yaml as yaml - with io.StringIO() as stream: - yaml.safe_dump(dict(self), stream) - filename.write(stream.getvalue()) - else: - raise NotImplementedError(filename.suffix) - - @classmethod - def load(cls, filename): - filename = path.Path(filename) - if filename.suffix == '.json': - return cls(json.loads(filename.read_text())) - elif filename.suffix in ('.yml', '.yaml'): - import ruamel.yaml as yaml - return cls(yaml.safe_load(filename.read_text())) - else: - raise NotImplementedError(filename.suffix) - - def __contains__(self, name): - try: - self[name] - return True - except KeyError: - return False - - def __getattr__(self, name): - if name.startswith('_'): - return super().__getattr__(name) - try: - return self[name] - except KeyError: - raise AttributeError(name) - - def __getitem__(self, name): - result = self._nested - for part in name.split(self.SEP): - try: - result = result[part] - except TypeError: - raise KeyError - if isinstance(result, dict): - result = type(self)(result) - return result - - def __setattr__(self, key, value): - if key.startswith('_'): - return super().__setattr__(key, value) - message = f"Tried to set key '{key}' on immutable config. Use update()." - raise AttributeError(message) - - def __setitem__(self, key, value): - if key.startswith('_'): - return super().__setitem__(key, value) - message = f"Tried to set key '{key}' on immutable config. Use update()." - raise AttributeError(message) - - def __reduce__(self): - return (type(self), (dict(self),)) - - def __str__(self): - lines = ['\nConfig:'] - keys, vals, typs = [], [], [] - for key, val in self.flat.items(): - keys.append(key + ':') - vals.append(self._format_value(val)) - typs.append(self._format_type(val)) - max_key = max(len(k) for k in keys) if keys else 0 - max_val = max(len(v) for v in vals) if vals else 0 - for key, val, typ in zip(keys, vals, typs): - key = key.ljust(max_key) - val = val.ljust(max_val) - lines.append(f'{key} {val} ({typ})') - return '\n'.join(lines) - - def update(self, *args, **kwargs): - result = self._flat.copy() - inputs = self._flatten(dict(*args, **kwargs)) - for key, new in inputs.items(): - if self.IS_PATTERN.match(key): - pattern = re.compile(key) - keys = {k for k in result if pattern.match(k)} - else: - keys = [key] - if not keys: - raise KeyError(f'Unknown key or pattern {key}.') - for key in keys: - old = result[key] - try: - if isinstance(old, int) and isinstance(new, float): - if float(int(new)) != new: - message = f"Cannot convert fractional float {new} to int." - raise ValueError(message) - result[key] = type(old)(new) - except (ValueError, TypeError): - raise TypeError( - f"Cannot convert '{new}' to type '{type(old).__name__}' " + - f"for key '{key}' with previous value '{old}'.") - return type(self)(result) - - def _flatten(self, mapping): - result = {} - for key, value in mapping.items(): - if isinstance(value, dict): - for k, v in self._flatten(value).items(): - if self.IS_PATTERN.match(key) or self.IS_PATTERN.match(k): - combined = f'{key}\\{self.SEP}{k}' - else: - combined = f'{key}{self.SEP}{k}' - result[combined] = v - else: - result[key] = value - return result - - def _nest(self, mapping): - result = {} - for key, value in mapping.items(): - parts = key.split(self.SEP) - node = result - for part in parts[:-1]: - if part not in node: - node[part] = {} - node = node[part] - node[parts[-1]] = value - return result - - def _ensure_keys(self, mapping): - for key in mapping: - assert not self.IS_PATTERN.match(key), key - return mapping - - def _ensure_values(self, mapping): - result = json.loads(json.dumps(mapping)) - for key, value in result.items(): - if isinstance(value, list): - value = list(value) - if isinstance(value, tuple): - if len(value) == 0: - message = 'Empty lists are disallowed because their type is unclear.' - raise TypeError(message) - if not isinstance(value[0], (str, float, int, bool)): - message = 'Lists can only contain strings, floats, ints, bools' - message += f' but not {type(value[0])}' - raise TypeError(message) - if not all(isinstance(x, type(value[0])) for x in value[1:]): - message = 'Elements of a list must all be of the same type.' - raise TypeError(message) - result[key] = value - return result - - def _format_value(self, value): - if isinstance(value, (list, tuple)): - return '[' + ', '.join(self._format_value(x) for x in value) + ']' - return str(value) - - def _format_type(self, value): - if isinstance(value, (list, tuple)): - assert len(value) > 0, value - return self._format_type(value[0]) + 's' - return str(type(value).__name__) diff --git a/conda-environment.yaml b/conda-environment.yaml index 753fa60..6651b6b 100644 --- a/conda-environment.yaml +++ b/conda-environment.yaml @@ -3,5 +3,5 @@ channels: - defaults - conda-forge dependencies: - - python=3.10 + - python=3.12 - pip diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..d6cce49 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,41 @@ +[build-system] +requires = ["setuptools>=77.0.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "bpo" +version = "0.1.0" +description = "Behaviour Policy Optimization: provably lower variance return estimates for off-policy reinforcement learning" +readme = "README.md" +requires-python = ">=3.12,<3.13" +dependencies = [ + "gymnasium[mujoco]>=1.1.1", + "tensorflow-cpu>=2.18.0,<2.19.0", + "tensorflow-probability==0.25.0", + "tf-keras>=2.18.0,<2.19.0", + "numpy==1.26.4", + "tensorboard>=2.18.0", + "tqdm>=4.66.5", + "jax==0.4.35", + "optax>=0.1.9", + "flax>=0.9.0", + "ml-dtypes>=0.4.1", + "ruamel.yaml>=0.18.0,<0.20.0", +] + +[project.scripts] +bpo = "bpo.run:main" + +[tool.setuptools] +include-package-data = true + +[tool.setuptools.packages.find] +include = ["bpo", "bpo.*"] + +[dependency-groups] +dev = ["pytest>=9.0.3"] +cuda12 = ["jax[cuda12]==0.4.35"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +pythonpath = ["."] diff --git a/requirements.txt b/requirements.txt index 472b11b..845504e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,13 +1,12 @@ gymnasium[mujoco]>=1.1.1 -tensorflow==2.17.0 -tensorflow-probability==0.22.1 -tensorstore==0.1.65 -tf-keras==2.17.0 +tensorflow-cpu>=2.18.0,<2.19.0 +tensorflow-probability==0.25.0 +tf-keras>=2.18.0,<2.19.0 numpy==1.26.4 -tensorboard>=2.17.1 -tqdm>=4.66.5 -jax==0.4.30 +tensorboard>=2.18.0 +tqdm>=4.66.5 +jax==0.4.35 optax>=0.1.9 flax>=0.9.0 ml-dtypes>=0.4.1 -ruamel.yaml<0.18.0 \ No newline at end of file +ruamel.yaml>=0.18.0,<0.20.0 diff --git a/run.py b/run.py index 15ec1ee..d12b7a3 100644 --- a/run.py +++ b/run.py @@ -1,257 +1,6 @@ -import warnings -warnings.filterwarnings("ignore", category=DeprecationWarning) -warnings.filterwarnings("ignore", category=UserWarning) +"""Launch training from the repository root.""" -import os - -from ppo import PPO -from bpo import BPO_PPO -from common.configs import Config, Flags, Path -from common.wrappers import TimeLimit, RewardMonitor, VecNormWrapper, NormWrapper, VecWrapper -from common.policies import PPOPolicy, BPOPolicy -from common.layers import NatureCNN - -import optax -from optax.schedules import linear_schedule -import flax.linen as nn -from flax.linen import initializers - -import importlib -import math -import ruamel.yaml as yaml -from typing import Dict, Any - -def make_env(env_id: str, max_episode_steps: int, **kwargs): - ctor = { - 'cartpole': 'envs.cartpole:DiscreteCartpole', - 'cont_cartpole': 'envs.cartpole:ContinuousCartpole', - 'ant': 'envs.mujoco:Ant', - 'half_cheetah': 'envs.mujoco:HalfCheetah', - 'hopper': 'envs.mujoco:Hopper', - 'walker_2d': 'envs.mujoco:Walker2D', - }[env_id] - if isinstance(ctor, str): - module, cls = ctor.split(':') - module = importlib.import_module(module) - ctor = getattr(module, cls) - env = ctor(**kwargs) - env = TimeLimit(env, max_episode_steps) - env = RewardMonitor(env) - return env - -def make_envs(n_envs: int, env_id: str, max_episode_steps: int, **kwargs): - ctor = { - 'cartpole': 'envs.cartpole:DiscreteCartpole', - 'cont_cartpole': 'envs.cartpole:ContinuousCartpole', - 'ant': 'envs.mujoco:Ant', - 'half_cheetah': 'envs.mujoco:HalfCheetah', - 'hopper': 'envs.mujoco:Hopper', - 'walker_2d': 'envs.mujoco:Walker2D', - }[env_id] - if isinstance(ctor, str): - module, cls = ctor.split(':') - module = importlib.import_module(module) - ctor = getattr(module, cls) - envs = [ - RewardMonitor( - TimeLimit( - ctor(**kwargs), max_episode_steps - ) - ) - for _ in range(n_envs) - ] - return envs - -def parse_policy_kwargs(config: Config): - - try: - algo_kwargs = Config(config[config.algo]) - except: - raise NotImplementedError(f"algorithm: {config.algo}") - - if config.bpo: - policy_class = BPOPolicy - else: - policy_class = PPOPolicy - - total_timesteps = config.run.total_timesteps - n_envs = config.run.n_envs - n_steps = algo_kwargs.n_steps - - if hasattr(algo_kwargs, "clip_range"): - if hasattr(algo_kwargs, "clip_decay") and algo_kwargs.clip_decay == "linear": - clip_schedule = linear_schedule(float(algo_kwargs.clip_range), 0.0, total_timesteps) - else: - clip_schedule = float(algo_kwargs.clip_range) - else: - clip_schedule = None - - if hasattr(algo_kwargs, "batch_size"): - batch_size = algo_kwargs.batch_size - else: - batch_size = 1 - - if hasattr(algo_kwargs, "n_epochs"): - n_epochs = algo_kwargs.n_epochs - else: - n_epochs = 1 - - decay_steps = (n_steps * n_envs // batch_size) * n_epochs \ - * math.ceil(total_timesteps / (n_steps * n_envs)) - - if algo_kwargs.use_featurizer: - featurizer_class = NatureCNN - else: - featurizer_class = None - - featurizer_kwargs = dict( - grayscale_obs=algo_kwargs.featurizer_kwargs.grayscale_obs, - normalize_images=algo_kwargs.featurizer_kwargs.normalize_images, - activation_fn={ - "tanh": nn.tanh, - "relu": nn.relu, - }[algo_kwargs.featurizer_kwargs.activation_fn], - kernel_init={ - "lecun_normal": initializers.lecun_normal(), - "orthogonal": initializers.orthogonal(), - }[algo_kwargs.featurizer_kwargs.kernel_init] - ) - - optimizer_class = { - "adam": optax.adam, - "rms_prop": optax.rmsprop - }[algo_kwargs.optimizer.opt] - - if algo_kwargs.optimizer.decay == "none": - lr_schedule = algo_kwargs.optimizer.learning_rate - elif algo_kwargs.optimizer.decay == "linear": - lr_schedule = linear_schedule(algo_kwargs.optimizer.learning_rate, 0.0, decay_steps) - else: - raise NotImplementedError(algo_kwargs.optimizer.decay) - - max_grad_norm = algo_kwargs.optimizer.max_grad_norm - - optimizer_kwargs = dict ( - eps=algo_kwargs.optimizer.eps - ) - - actor_kwargs = dict( - n_units=algo_kwargs.actor_kwargs.n_units, - log_std_init=algo_kwargs.actor_kwargs.log_std_init, - activation_fn={ - "tanh": nn.tanh, - "relu": nn.relu, - }[algo_kwargs.actor_kwargs.activation_fn], - kernel_init={ - "lecun_normal": initializers.lecun_normal(), - "orthogonal": initializers.orthogonal(), - }[algo_kwargs.actor_kwargs.kernel_init] - ) - - critic_kwargs = dict( - n_units=algo_kwargs.critic_kwargs.n_units, - use_layer_norm=algo_kwargs.critic_kwargs.use_layer_norm, - activation_fn={ - "tanh": nn.tanh, - "relu": nn.relu, - }[algo_kwargs.critic_kwargs.activation_fn], - kernel_init={ - "lecun_normal": initializers.lecun_normal(), - "orthogonal": initializers.orthogonal(), - }[algo_kwargs.critic_kwargs.kernel_init] - ) - - q_vf_kwargs = dict( - n_units=algo_kwargs.q_vf_kwargs.n_units, - use_layer_norm=algo_kwargs.q_vf_kwargs.use_layer_norm, - use_zero_norm_final=algo_kwargs.q_vf_kwargs.use_zero_norm_final, - activation_fn={ - "tanh": nn.tanh, - "relu": nn.relu, - }[algo_kwargs.q_vf_kwargs.activation_fn], - kernel_init={ - "lecun_normal": initializers.lecun_normal(), - "orthogonal": initializers.orthogonal(), - }[algo_kwargs.q_vf_kwargs.kernel_init] - ) - - policy_kwargs = dict( - featurizer_class=featurizer_class, - featurizer_kwargs=featurizer_kwargs, - optimizer_class=optimizer_class, - optimizer_kwargs=optimizer_kwargs, - actor_kwargs=actor_kwargs, - critic_kwargs=critic_kwargs, - q_vf_kwargs=q_vf_kwargs, - ) - - algo_kwargs = {**algo_kwargs} - algo_kwargs.update({ - "learning_rate": lr_schedule, - "clip_range": clip_schedule, - "max_grad_norm": max_grad_norm, - "policy_class": policy_class, - "policy_kwargs": policy_kwargs, - }) - - return algo_kwargs - -def main(argv=None): - parsed, other = Flags(configs=['defaults']).parse_known(argv) - configs = yaml.YAML(typ='safe').load( - (Path(__file__).parent / 'configs.yaml').read()) - config = Config(configs['defaults']) - for name in parsed.configs: - config = config.update(configs[name]) - config = Flags(config).parse(other) - - print(config) - - try: - algo_kwargs = Config(config[config.algo]) - except: - raise NotImplementedError(f"algorithm: {config.algo}") - - if config.run.n_envs > 1: - envs = make_envs(config.run.n_envs, config.env.env_id, config.env.max_episode_steps) - if config.env.normalize_obs or config.env.normalize_rew: - env = VecNormWrapper(envs, norm_obs=config.env.normalize_obs, norm_rew=config.env.normalize_rew, training=True, gamma=algo_kwargs.gamma) - else: - env = VecWrapper(envs) - else: - env = make_env(config.env.env_id, config.env.max_episode_steps) - if config.env.normalize_obs or config.env.normalize_rew: - env = NormWrapper(env, norm_obs=config.env.normalize_obs, norm_rew=config.env.normalize_rew, training=True, gamma=algo_kwargs.gamma) - - eval_env = make_env(config.env.env_id, config.env.max_episode_steps) - - base_kwargs = dict( - tensorboard_logdir=config.run.logdir, - seed=config.run.seed, - device=config.run.device, - verbose=config.run.verbose, - eval_env=eval_env, - ) - - algo_id = f"{'bpo_' if bool(config.bpo) else ''}{config.algo}" - - algo_class = { - "ppo": PPO, - "bpo_ppo": BPO_PPO, - }[algo_id] - - algo_kwargs = parse_policy_kwargs(config) - merged = {**base_kwargs, **algo_kwargs} - algo = algo_class(env, **merged) - - run_kwargs = dict( - num_eval_episodes=config.run.eval_episodes, - eval_freq=config.run.eval_every, - log_freq=config.run.log_every, - stats_window_size=config.run.stats_window_size, - ) - - algo.train(config.run.total_timesteps, **run_kwargs) +from bpo.run import main if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/scripts/check_requirements_sync.py b/scripts/check_requirements_sync.py new file mode 100644 index 0000000..37e8e6d --- /dev/null +++ b/scripts/check_requirements_sync.py @@ -0,0 +1,124 @@ +"""Check that requirements.txt and pyproject.toml stay in lockstep. + +The conda installation flow relies on ``requirements.txt`` while uv/PyPI +consume ``[project.dependencies]`` in ``pyproject.toml``. This script +compares the two with strict string equality so the files cannot drift. +""" + +from __future__ import annotations + +import re +import sys +from pathlib import Path + +try: + import tomllib +except ModuleNotFoundError: # Python < 3.11 + import tomli as tomllib # type: ignore[no-redefine] + +ROOT = Path(__file__).resolve().parents[1] +PYPROJECT = ROOT / "pyproject.toml" +REQUIREMENTS = ROOT / "requirements.txt" + +# PEP 503 name normalization: lowercase, collapse runs of -, _ and . to a +# single -. +_NAME_PREFIX = re.compile(r"^([A-Za-z0-9][A-Za-z0-9._-]*)(.*)$", re.DOTALL) +_NORMALIZE = re.compile(r"[-_.]+") + + +def _split_requirement(requirement: str) -> tuple[str, str]: + """Return (canonical_name, normalized_remainder) for a requirement string. + + The remainder (extras, version specifiers, environment markers) has all + whitespace stripped so comparison is purely on content, not formatting. + """ + match = _NAME_PREFIX.match(requirement.strip()) + if not match: + raise ValueError(f"Could not parse requirement: {requirement!r}") + name, remainder = match.groups() + canonical = _NORMALIZE.sub("-", name).lower() + normalized_remainder = re.sub(r"\s+", "", remainder) + return canonical, normalized_remainder + + +def parse_pyproject(path: Path) -> dict[str, str]: + with path.open("rb") as handle: + data = tomllib.load(handle) + dependencies = data.get("project", {}).get("dependencies", []) + result: dict[str, str] = {} + for raw in dependencies: + name, remainder = _split_requirement(raw) + if name in result: + raise SystemExit( + f"Duplicate dependency '{name}' in {path.name} [project.dependencies]" + ) + result[name] = remainder + return result + + +def parse_requirements(path: Path) -> dict[str, str]: + result: dict[str, str] = {} + duplicates: list[str] = [] + for lineno, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): + stripped = line.split("#", 1)[0].strip() + if not stripped: + continue + name, remainder = _split_requirement(stripped) + if name in result: + duplicates.append(f" line {lineno}: {name}") + continue + result[name] = remainder + if duplicates: + raise SystemExit( + f"Duplicate dependencies in {path.name}:\n" + "\n".join(duplicates) + ) + return result + + +def main() -> int: + pyproject = parse_pyproject(PYPROJECT) + requirements = parse_requirements(REQUIREMENTS) + + missing = sorted(set(pyproject) - set(requirements)) + extra = sorted(set(requirements) - set(pyproject)) + mismatched = sorted( + name + for name in set(pyproject) & set(requirements) + if pyproject[name] != requirements[name] + ) + + if not (missing or extra or mismatched): + print( + f"OK: {REQUIREMENTS.name} and {PYPROJECT.name} " + f"[project.dependencies] are in sync ({len(pyproject)} dependencies)." + ) + return 0 + + print("requirements.txt and pyproject.toml [project.dependencies] disagree:\n") + if missing: + print("Missing from requirements.txt (present in pyproject.toml):") + for name in missing: + print(f" - {name}{pyproject[name]}") + print() + if extra: + print("Extra in requirements.txt (absent from pyproject.toml):") + for name in extra: + print(f" - {name}{requirements[name]}") + print() + if mismatched: + print("Specifier mismatches:") + for name in mismatched: + print( + f" - {name}: pyproject '{name}{pyproject[name]}' " + f"!= requirements '{name}{requirements[name]}'" + ) + print() + print( + "Update requirements.txt and [project.dependencies] so they match " + "exactly (base dependencies only)." + ) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_gymnasium_envs.py b/tests/test_gymnasium_envs.py new file mode 100644 index 0000000..5c522e1 --- /dev/null +++ b/tests/test_gymnasium_envs.py @@ -0,0 +1,78 @@ +"""Gymnasium contract checks for the custom cartpole environments. + +Env classes are imported directly so these tests do not pull in the +training stack. +""" + +import numpy as np +import pytest + +import gymnasium +from bpo.envs.cartpole import ContinuousCartpole, DiscreteCartpole + +ENV_CLASSES = [ + DiscreteCartpole, + ContinuousCartpole, +] + + +def _make(cls, seed=None): + env = cls() + obs, info = env.reset(seed=seed) + return env, obs, info + + +def _action(env): + if isinstance(env.action_space, gymnasium.spaces.Discrete): + action = 0 + assert isinstance(action, int) + assert action in (0, 1) + return action + return np.zeros(env.action_space.shape, dtype=env.action_space.dtype) + + +@pytest.mark.parametrize("cls", ENV_CLASSES, ids=lambda cls: cls.__name__) +def test_spaces_are_gymnasium_spaces(cls): + env, _, _ = _make(cls, seed=123) + assert isinstance(env.observation_space, gymnasium.spaces.Space) + assert isinstance(env.action_space, gymnasium.spaces.Space) + assert isinstance(env.observation_space, gymnasium.spaces.Box) + if cls is DiscreteCartpole: + assert isinstance(env.action_space, gymnasium.spaces.Discrete) + assert env.action_space.n == 2 + else: + assert isinstance(env.action_space, gymnasium.spaces.Box) + + +@pytest.mark.parametrize("cls", ENV_CLASSES, ids=lambda cls: cls.__name__) +def test_reset_returns_obs_and_info(cls): + env, obs, info = _make(cls, seed=123) + assert isinstance(info, dict) + assert isinstance(obs, np.ndarray) + assert obs.shape == env.observation_space.shape + + +@pytest.mark.parametrize("cls", ENV_CLASSES, ids=lambda cls: cls.__name__) +def test_reset_seed_is_reproducible(cls): + env_a, obs_a, _ = _make(cls, seed=123) + env_b, obs_b, _ = _make(cls, seed=123) + np.testing.assert_allclose(obs_a, obs_b) + + +@pytest.mark.parametrize("cls", ENV_CLASSES, ids=lambda cls: cls.__name__) +def test_step_returns_gymnasium_tuple(cls): + env, obs, _ = _make(cls, seed=123) + action = _action(env) + assert env.action_space.contains(action) + + result = env.step(action) + assert len(result) == 5 + next_obs, reward, terminated, truncated, info = result + + assert isinstance(next_obs, np.ndarray) + assert next_obs.shape == env.observation_space.shape + assert np.ndim(reward) == 0 + assert isinstance(terminated, (bool, np.bool_)) + assert isinstance(truncated, (bool, np.bool_)) + assert isinstance(info, dict) + assert obs.shape == env.observation_space.shape diff --git a/tests/test_imports.py b/tests/test_imports.py new file mode 100644 index 0000000..53bc74e --- /dev/null +++ b/tests/test_imports.py @@ -0,0 +1,25 @@ +"""Import surface for the pinned scientific stack. + +These stay independent of the training code so a broken dependency install +fails here rather than inside an algorithm test. +""" + +import importlib + +import pytest + +REQUIRED_MODULES = [ + "numpy", + "jax", + "jaxlib", + "tensorflow", + "tensorflow_probability", + "flax", + "optax", + "gymnasium", +] + + +@pytest.mark.parametrize("module_name", REQUIRED_MODULES) +def test_required_module_imports(module_name): + importlib.import_module(module_name) diff --git a/tests/test_smoke.py b/tests/test_smoke.py new file mode 100644 index 0000000..165de0d --- /dev/null +++ b/tests/test_smoke.py @@ -0,0 +1,115 @@ +"""Smoke coverage for environment construction and training launch. + +Imports of the training stack happen inside each test so a broken install +shows up as a test failure rather than a collection error. +""" + +import numpy as np +import pytest + +import gymnasium + +ENV_IDS = [ + "cartpole", + "cont_cartpole", + "ant", + "half_cheetah", + "hopper", + "walker_2d", +] + + +def _action(env): + if isinstance(env.action_space, gymnasium.spaces.Discrete): + return 0 + return np.zeros(env.action_space.shape, dtype=env.action_space.dtype) + + +def _lean_argv(logdir, *extra): + """One rollout and one update. Batch sizes match the rollout length. + + BPOReplayBuffer.get asserts the minibatch is no larger than the number + of stored transitions, and the rollout buffer is full only at n_steps. + """ + return [ + *extra, + "--run.total_timesteps", + "4", + "--run.n_envs", + "1", + "--run.eval_every", + "1000000", + "--run.log_every", + "1000000", + "--run.stats_window_size", + "1", + "--run.logdir", + str(logdir), + "--env.max_episode_steps", + "20", + "--ppo.n_steps", + "4", + "--ppo.batch_size", + "4", + "--ppo.n_epochs", + "1", + "--ppo.n_vf_epochs", + "1", + "--ppo.n_mu_epochs", + "1", + "--ppo.vf_batch_size", + "4", + "--ppo.mu_batch_size", + "4", + "--ppo.replay_size", + "32", + ] + + +@pytest.mark.parametrize("env_id", ENV_IDS) +def test_make_env(env_id): + from bpo.run import make_env + + env = make_env(env_id, max_episode_steps=5) + obs, info = env.reset(seed=1) + assert isinstance(info, dict) + assert isinstance(obs, np.ndarray) + assert obs.shape == env.observation_space.shape + + action = _action(env) + assert env.action_space.contains(action) + result = env.step(action) + assert len(result) == 5 + next_obs, reward, terminated, truncated, step_info = result + assert isinstance(next_obs, np.ndarray) + assert next_obs.shape == env.observation_space.shape + assert np.ndim(reward) == 0 + assert isinstance(terminated, (bool, np.bool_)) + assert isinstance(truncated, (bool, np.bool_)) + assert isinstance(step_info, dict) + + +def test_training_launches_ppo(tmp_path): + from bpo.run import main + + main(_lean_argv(tmp_path, "--configs", "cartpole_ppo")) + + +def test_training_launches_bpo(tmp_path): + from bpo.run import main + + main(_lean_argv(tmp_path, "--configs", "cartpole_ppo", "ppo_bpo_zero")) + + +def test_training_launches_ppo_continuous(tmp_path): + from bpo.run import main + + main( + _lean_argv( + tmp_path, + "--configs", + "cartpole_ppo", + "--env.env_id", + "cont_cartpole", + ) + ) diff --git a/uv.lock b/uv.lock new file mode 100644 index 0000000..dc148d9 --- /dev/null +++ b/uv.lock @@ -0,0 +1,1351 @@ +version = 1 +revision = 5 +requires-python = "==3.12.*" + +[[package]] +name = "absl-py" +version = "2.5.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1f/1d/58e2b5a6e4d703ccb2a029943d665974cb3d5a4fb2b3e3675dd03a9df10e/absl_py-2.5.1.tar.gz", hash = "sha256:286e71c82c1a38e75bbcf185f9b37d0305ad7786535107cb49bf4df9ff2e1f95", size = 118387, upload-time = "2026-10-09T08:38:58.037Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/7d/01e62f59e4166af1be6238d9f2b7f3453630c51b8044b412d05b5606532a/absl_py-2.5.1-py3-none-any.whl", hash = "sha256:721200f2f0e9960f2ca9dc3a2a706b201f5f75d812c158f059cbbe29eeafbdb8", size = 137673, upload-time = "2026-10-09T08:38:56.629Z" }, +] + +[[package]] +name = "astunparse" +version = "1.6.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "six" }, + { name = "wheel" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f3/af/4182184d3c338792894f34a62672919db7ca008c89abee9b564dd34d8029/astunparse-1.6.3.tar.gz", hash = "sha256:5ad93a8456f0d084c3456d059fd9a92cce667963232cbf763eac3bc5b7940872", size = 18290, upload-time = "2019-12-22T18:12:13.129Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2b/03/13dde6512ad7b4557eb792fbcf0c653af6076b81e5941d36ec61f7ce6028/astunparse-1.6.3-py2.py3-none-any.whl", hash = "sha256:c2652417f2c8b5bb325c885ae329bdf3f86424075c4fd1a128674bc6fba4b8e8", size = 12732, upload-time = "2019-12-22T18:12:11.297Z" }, +] + +[[package]] +name = "attrs" +version = "26.1.0" +source = { registry = "https://pypi.org/simple" } +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/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]] +name = "bpo" +version = "0.1.0" +source = { editable = "." } +dependencies = [ + { name = "flax" }, + { name = "gymnasium", extra = ["mujoco"] }, + { name = "jax" }, + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "optax" }, + { name = "ruamel-yaml" }, + { name = "tensorboard" }, + { name = "tensorflow-cpu" }, + { name = "tensorflow-probability" }, + { name = "tf-keras" }, + { name = "tqdm" }, +] + +[package.dev-dependencies] +cuda12 = [ + { name = "jax", extra = ["cuda12"] }, +] +dev = [ + { name = "pytest" }, +] + +[package.metadata] +requires-dist = [ + { name = "flax", specifier = ">=0.9.0" }, + { name = "gymnasium", extras = ["mujoco"], specifier = ">=1.1.1" }, + { name = "jax", specifier = "==0.4.35" }, + { name = "ml-dtypes", specifier = ">=0.4.1" }, + { name = "numpy", specifier = "==1.26.4" }, + { name = "optax", specifier = ">=0.1.9" }, + { name = "ruamel-yaml", specifier = ">=0.18.0,<0.20.0" }, + { name = "tensorboard", specifier = ">=2.18.0" }, + { name = "tensorflow-cpu", specifier = ">=2.18.0,<2.19.0" }, + { name = "tensorflow-probability", specifier = "==0.25.0" }, + { name = "tf-keras", specifier = ">=2.18.0,<2.19.0" }, + { name = "tqdm", specifier = ">=4.66.5" }, +] + +[package.metadata.requires-dev] +cuda12 = [{ name = "jax", extras = ["cuda12"], specifier = "==0.4.35" }] +dev = [{ name = "pytest", specifier = ">=9.0.3" }] + +[[package]] +name = "certifi" +version = "2026.10.23" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/11/b524cbd9bbdaca58608e4cf4053fd463417549e1959fa8d3aa18319f9555/certifi-2026.10.23.tar.gz", hash = "sha256:a9895e6e2758b371e0545e321a3a8fd1fcfd9829ada8824a936dfb16aaaf36cd", size = 136225, upload-time = "2026-10-23T11:36:16.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/32/d1/a3652fe2c19134d0d15922177821d863cbd7e842d53a21bbe7e8e3a1b5ac/certifi-2026.10.23-py3-none-any.whl", hash = "sha256:e9b54de76e530bb6cf0ccd036787103acb58fd42ba930ea8bf948ddbaa2f61d1", size = 135082, upload-time = "2026-10-23T11:36:14.763Z" }, +] + +[[package]] +name = "charset-normalizer" +version = "3.5.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/33/1c/f41d4e74c28ab327ff3acd36053f7ea506c55872d7a90b0fa71aa3ab0c89/charset_normalizer-3.5.2.tar.gz", hash = "sha256:39de2a259fc954455c57274dc94c79d5842774e1247a016aff30bc0efed0f4ef", size = 172659, upload-time = "2026-09-30T04:39:23.398Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/c8/693809898870237d82785a03f3b2b58fe4c9f14669f84a7d4e623c92a59e/charset_normalizer-3.5.2-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:ed2a239c0ea213acc1908150a3037257083c7c083128f1a4cec2ec4b97dca491", size = 367780, upload-time = "2026-09-30T04:35:30.888Z" }, + { url = "https://files.pythonhosted.org/packages/c9/87/2fea8c13dc24b3ca9c6f803a5b2dfdeae73eb4f9e12c7885ed908ff0433c/charset_normalizer-3.5.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b91363207bd9dc966a691e959bb47f64b30f7ac4b072be9968b366982f7db77c", size = 246730, upload-time = "2026-09-30T04:35:32.286Z" }, + { url = "https://files.pythonhosted.org/packages/a8/9e/09efac30b937722f46d3110ba30b875b24b2e3a266ed746cc4e376a94d80/charset_normalizer-3.5.2-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:38a873987f3be698494da8b2e3085e29da02da7b633dce73e79c699a113d7bf0", size = 237707, upload-time = "2026-09-30T04:35:33.709Z" }, + { url = "https://files.pythonhosted.org/packages/9e/18/70d76670b13686237863a379928d60bd10e021f17d243ab3d7014c4a5f4e/charset_normalizer-3.5.2-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:355ad8011081dec5412240c087a9a0c9d4d5039f3ed11a3f13e18c2b29b56c51", size = 273050, upload-time = "2026-09-30T04:35:35.138Z" }, + { url = "https://files.pythonhosted.org/packages/54/e2/77a8b09d5adc013ed07b95b01b8b8fa5441c4e810e83ee7e4aae2fa4d91a/charset_normalizer-3.5.2-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:ee21e28f0430bd6dc9086c6e525d5e818a44a5ad19720c8a0ef766792f3eb5e5", size = 270345, upload-time = "2026-09-30T04:35:36.502Z" }, + { url = "https://files.pythonhosted.org/packages/7f/c5/38806a25ab5e65fc178f39affeda20858efafede2fce1ffc2556cfc9fe73/charset_normalizer-3.5.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3d31298449090ab8d47b7b1b2a555ff73cac7ed438a08b7ac160980c7ebed649", size = 257601, upload-time = "2026-09-30T04:35:37.919Z" }, + { url = "https://files.pythonhosted.org/packages/ae/8d/213565184708fdb263ae55e2c04ee1ff748129dd65d48ed0e3502da9c85a/charset_normalizer-3.5.2-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:5cde776b7cc66e4f6c99612cea4aa7269aa65863f7a15841b2c264f103822f4e", size = 252222, upload-time = "2026-09-30T04:35:39.544Z" }, + { url = "https://files.pythonhosted.org/packages/7e/24/76d2cefc25472531e4c5c7dfff68865eb1c39b78482f0fdc15b46f047830/charset_normalizer-3.5.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ae4f5fea5b8b8ccff88238cc8569303e5ee95efae67fa62922a311397a71f346", size = 248482, upload-time = "2026-09-30T04:35:41.088Z" }, + { url = "https://files.pythonhosted.org/packages/7d/dc/65a801b66ab4c197e22c433ab25e7ac24324ac6f45a2269aca42cce309bf/charset_normalizer-3.5.2-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:f7d486c83842422badd511868fd8a9a20e9407ace71564b6af47ce7e60a336c1", size = 241206, upload-time = "2026-09-30T04:35:42.59Z" }, + { url = "https://files.pythonhosted.org/packages/a7/95/ca9b5eabde673002c6f1e7ada1b223916fe18f6d661da7aabd4d643718f1/charset_normalizer-3.5.2-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:11a4d68a6ecda3292cb1e50239e111543ba5d709bb62a6b4ea1afcfa729d8875", size = 273190, upload-time = "2026-09-30T04:35:44.347Z" }, + { url = "https://files.pythonhosted.org/packages/2d/8b/803b4d2a3f6e1740f63f1e87b04d14b42f3d4fdfe6ed7d4db2d34102b14f/charset_normalizer-3.5.2-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:d6734d2ef8a50fbf8445c139477da401f50d62a0606bf00e20ec6d87773fefb1", size = 253527, upload-time = "2026-09-30T04:35:45.915Z" }, + { url = "https://files.pythonhosted.org/packages/a9/55/93c0e5dbd085ae0471346026abbe7e0db9ea2d6fea74e51f0b5a46f233a7/charset_normalizer-3.5.2-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:a815775b6c38d4e0ff7bcffbeba67feded90202bb6a226b8dd35f1c855217413", size = 271285, upload-time = "2026-09-30T04:35:47.49Z" }, + { url = "https://files.pythonhosted.org/packages/95/69/0dbd0e0b9b16cfa816cdfcb3e2e3854a1f680dc07fb1245ea125e7448060/charset_normalizer-3.5.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:23851fb4e1b85ed3f6c2a27b777cdfe2e19fb5b38429a8faf38c7542b7665869", size = 260010, upload-time = "2026-09-30T04:35:48.996Z" }, + { url = "https://files.pythonhosted.org/packages/58/9d/e7b88e7b1bf403590c3b573277b5e1e488c68c7a6fbacca310a2c324e90c/charset_normalizer-3.5.2-cp312-cp312-win32.whl", hash = "sha256:db19d07e2e0129e974a0e65d0064fc222a446cd5122c2fd4184d2af9fc734a9e", size = 184125, upload-time = "2026-09-30T04:35:50.777Z" }, + { url = "https://files.pythonhosted.org/packages/eb/e6/e6e083884cbcfd49c64865af05027fe7011be7b2d9179524f099a1b611f3/charset_normalizer-3.5.2-cp312-cp312-win_amd64.whl", hash = "sha256:780fbe7cab297b81dad9fb8dc5eb003c0468ffb0d9e5f65068c53a34661a96bc", size = 207486, upload-time = "2026-09-30T04:35:52.194Z" }, + { url = "https://files.pythonhosted.org/packages/c4/e3/017aea0911ada7405a825c7d937eb3a13009664e2f5b38e8c4bbf2abf894/charset_normalizer-3.5.2-cp312-cp312-win_arm64.whl", hash = "sha256:e2af3aad578aa6bd1384bcf4750fc285e5a9de53f40b7d41e5a0bf748edeb2b3", size = 196734, upload-time = "2026-09-30T04:35:53.636Z" }, + { url = "https://files.pythonhosted.org/packages/8c/ab/176fbfd5b64939c55d652366aa5b9ef1d767af207a3aa6ebeb0d226c484d/charset_normalizer-3.5.2-cp37-abi3-macosx_10_9_universal2.whl", hash = "sha256:4275811936e2f06feff5e598fb42a1b7ae852da8e39605211892b56b81a34efd", size = 331815, upload-time = "2026-09-30T04:38:26.216Z" }, + { url = "https://files.pythonhosted.org/packages/7e/84/371eac6b30bdbcbf2d632a1a01809103459216fcaae61b8b8d922c1bfb8a/charset_normalizer-3.5.2-cp37-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:1c50fe28bbc2ced33386f298650d91218076c05420e6cbd790b913adc41659e7", size = 253276, upload-time = "2026-09-30T04:38:28.032Z" }, + { url = "https://files.pythonhosted.org/packages/43/6f/c4fbae58febff71709c51bc7e18fdfa55341dc382704740f9f0cbf03817b/charset_normalizer-3.5.2-cp37-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d19fbd981a488e22cd04883659ca6b08f50b5974f9fd7c95655ef6a043e5893f", size = 241239, upload-time = "2026-09-30T04:38:29.732Z" }, + { url = "https://files.pythonhosted.org/packages/61/71/458c3f42164a07d0c5210798e9e704b39e540a6793b05aba67f3a35243a9/charset_normalizer-3.5.2-cp37-abi3-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:0fed1d06615f022ee3b13caf5e8b180cfea32bb2c5aded8a9d44277afc040f93", size = 231121, upload-time = "2026-09-30T04:38:31.462Z" }, + { url = "https://files.pythonhosted.org/packages/09/54/ab9e89367076f6331bb6c65c4bf14a5361fa5191cb6561bf534f18504e1b/charset_normalizer-3.5.2-cp37-abi3-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:838dcc90063569a0448120554591a1d6c4a4ffe11babf048908793154ab86ade", size = 260350, upload-time = "2026-09-30T04:38:33.239Z" }, + { url = "https://files.pythonhosted.org/packages/7c/c1/061431ecc688d9d76602502cb57cc01e691e682c18f1beb45f9673b5bbd2/charset_normalizer-3.5.2-cp37-abi3-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2ce45c6627b22c47e390bc91a41c3d13032192e699fa0bea96e9671b373d69b0", size = 255430, upload-time = "2026-09-30T04:38:34.865Z" }, + { url = "https://files.pythonhosted.org/packages/8d/1f/20c8949f0676f7ab811abdeb7f4d7f1cbc6e61ff20bef08b44edeb092bc8/charset_normalizer-3.5.2-cp37-abi3-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:0774bf9bf620249fee3e0b8b9fd3065de213be30f3aa94ce2494b3b638949e26", size = 250612, upload-time = "2026-09-30T04:38:36.649Z" }, + { url = "https://files.pythonhosted.org/packages/2b/9e/46f2fa4c431fc98c4ae76a8cb5bdca54e0341e3cfc3fcfd8e82740250818/charset_normalizer-3.5.2-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:1db38f4c5496827c1a501846d64d14c3b80c7e6714e406cd7dc36a9899fa1011", size = 242083, upload-time = "2026-09-30T04:38:38.26Z" }, + { url = "https://files.pythonhosted.org/packages/bd/39/559be29a0c0f086e0bba6922babd38916cc5e0b58ced4de13ee01ea05508/charset_normalizer-3.5.2-cp37-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:304d8e4d493af723536393eee0c689eb7813f4a474c8b479dee63f1fdd98f621", size = 232738, upload-time = "2026-09-30T04:38:39.81Z" }, + { url = "https://files.pythonhosted.org/packages/ff/6c/387b0e4f756a282831c1d9fc6aeb6c51ca4507ca202767c8de15ce9b12e2/charset_normalizer-3.5.2-cp37-abi3-musllinux_1_2_ppc64le.whl", hash = "sha256:9b7f416ff0978e2f2249330527f0ad6fa02f4932e6199692d3b52da2048c19e4", size = 260703, upload-time = "2026-09-30T04:38:41.346Z" }, + { url = "https://files.pythonhosted.org/packages/96/92/1fdf015f09ef449f50d3ac4b67c90887c9c318b727daa95cc4f866e6521d/charset_normalizer-3.5.2-cp37-abi3-musllinux_1_2_riscv64.whl", hash = "sha256:01077390b03f7988f11d700a2194e69b119741a86b1a638b1db88891e3eced8e", size = 247622, upload-time = "2026-09-30T04:38:42.937Z" }, + { url = "https://files.pythonhosted.org/packages/dc/3c/8e7b8a5671ad5d433669fb2a76f1a0164df2d9b1718b0206bc2a16d840cc/charset_normalizer-3.5.2-cp37-abi3-musllinux_1_2_s390x.whl", hash = "sha256:7e841fb9010836c992c9f12fcbd43a831de93a5f726fc1ccd8ca1d0268c5014c", size = 257500, upload-time = "2026-09-30T04:38:44.604Z" }, + { url = "https://files.pythonhosted.org/packages/b4/f0/45b579df5cabc1d5d53ea1cc35e8437d3ca768c0acccc7041517cb6fbb32/charset_normalizer-3.5.2-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:9cae88599c7219005d879f98e5ed53341e9a122af585e1091200358a3003d2a0", size = 255100, upload-time = "2026-09-30T04:38:46.289Z" }, + { url = "https://files.pythonhosted.org/packages/31/68/fdec18a343f5fb3f310588dd478b09ac4799e0b187dbade3a8cd776f03ef/charset_normalizer-3.5.2-cp37-abi3-win32.whl", hash = "sha256:01b0c0d2262a9e28e8484a278c7e1b5d650e3ac8cf2683d2967e25899f208bdf", size = 174499, upload-time = "2026-09-30T04:38:47.999Z" }, + { url = "https://files.pythonhosted.org/packages/9d/8a/b618149cc5207943a0242068d7a27897f56a62947b5a039085f2a22029f8/charset_normalizer-3.5.2-cp37-abi3-win_amd64.whl", hash = "sha256:9f56f72050826f63dcee7a7f55b0a77168cb3bfc553fd405e7f8f9ece75a4036", size = 200092, upload-time = "2026-09-30T04:38:49.707Z" }, + { url = "https://files.pythonhosted.org/packages/03/cf/4c66866fa9e2b1c78e3c911516d1de497a677b7ac60f1eceda74ce777ca3/charset_normalizer-3.5.2-cp37-abi3-win_arm64.whl", hash = "sha256:40ab6bffa02ae10a0581e6c198be7d2d8ca5c2a0c64e4ed3465d766df457573e", size = 294363, upload-time = "2026-09-30T04:38:51.312Z" }, + { url = "https://files.pythonhosted.org/packages/fc/ad/d07d7862a62ffa6d79d68074d14823243dd235a77c45262acbf6adeb28bf/charset_normalizer-3.5.2-py3-none-any.whl", hash = "sha256:b6b751274acb69d77b3323d6b7dbaa3c7fdfc1eb829b7eb61d262f32e1af9685", size = 68872, upload-time = "2026-09-30T04:39:21.828Z" }, +] + +[[package]] +name = "chex" +version = "0.1.90" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "jax" }, + { name = "jaxlib" }, + { name = "numpy" }, + { name = "setuptools" }, + { name = "toolz" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/77/70/53c7d404ce9e2a94009aea7f77ef6e392f6740e071c62683a506647c520f/chex-0.1.90.tar.gz", hash = "sha256:d3c375aeb6154b08f1cccd2bee4ed83659ee2198a6acf1160d2fe2e4a6c87b5c", size = 92363, upload-time = "2025-07-23T19:50:47.945Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6f/3d/46bb04776c465cea2dd8aa2d4b61ab610b707f798f47838ef7e6105b025c/chex-0.1.90-py3-none-any.whl", hash = "sha256:fce3de82588f72d4796e545e574a433aa29229cbdcf792555e41bead24b704ae", size = 101047, upload-time = "2025-07-23T19:50:46.603Z" }, +] + +[[package]] +name = "cloudpickle" +version = "3.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, +] + +[[package]] +name = "colorama" +version = "0.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, +] + +[[package]] +name = "decorator" +version = "5.3.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/60/8b/32f9823da46cde7df2087faa08cd98d01b908f8dcab982cdba9c84e85355/decorator-5.3.1.tar.gz", hash = "sha256:4cbcdd55a6efadb9dbea26b858f4fb3264567b52d69ca0d25b721b553f60ea82", size = 58084, upload-time = "2026-05-18T06:03:28.057Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/05/7f/798705f5296a58ca505d600456748d1be48078eac8a7050d8a98bc9edb89/decorator-5.3.1-py3-none-any.whl", hash = "sha256:f47fe6fdbd2edd623ecfe36875d37aba411624e2670dd395dddae1358689bb3c", size = 10365, upload-time = "2026-05-18T06:03:26.517Z" }, +] + +[[package]] +name = "dm-tree" +version = "0.1.10" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "attrs" }, + { name = "numpy" }, + { name = "wrapt" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5a/66/a3ec619d22b6baffa5ab853e8dc6ec9d0c837127948af59bb15b988d7312/dm_tree-0.1.10.tar.gz", hash = "sha256:22f37b599e01cc3402a17f79c257a802aebd8d326de05b54657650845956208a", size = 35748, upload-time = "2026-03-31T17:35:39.03Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/34/a1/17e0d68eec978c483db4712b14d083ee01484381b29ea85edb2b20210bd0/dm_tree-0.1.10-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:94af18e4fd22ce69eccae89eeed8ed498b6b4cc4957f4ed10b4160e59f620e1d", size = 315976, upload-time = "2026-03-31T17:35:15.21Z" }, + { url = "https://files.pythonhosted.org/packages/cc/6f/ed603715fbc29c887a8985252e2cfe0d449497aea96bac51010159771617/dm_tree-0.1.10-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b442a0c1e9d0960e0314a2e4af81fd328a87921b6d6db6dc41bfa420536884d6", size = 184053, upload-time = "2026-03-31T17:35:16.512Z" }, + { url = "https://files.pythonhosted.org/packages/83/eb/1d55c679cee9a54e552480d308535753c72e2250cf720d7aa777bff2a4fe/dm_tree-0.1.10-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:012c2b376e88d3685c73a4b5c23be41fe933e14e380dcd90172971690b0e02d2", size = 186506, upload-time = "2026-03-31T17:35:17.593Z" }, + { url = "https://files.pythonhosted.org/packages/89/2d/adef6924f8dc7f1665eea4ce066387820c14a629d0e1005568892d56ea6a/dm_tree-0.1.10-cp312-cp312-win_amd64.whl", hash = "sha256:da8d5b8995bea1b6bb93f457e0dad5d16e6e2344a6488ced55320e7f3fd50f56", size = 112708, upload-time = "2026-03-31T17:35:18.699Z" }, +] + +[[package]] +name = "etils" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +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/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 = "typing-extensions" }, + { name = "zipp" }, +] +epy = [ + { name = "typing-extensions" }, +] + +[[package]] +name = "farama-notifications" +version = "0.0.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ec/91/14397890dde30adc4bee6462158933806207bc5dd10d7b4d09d5c33845cf/farama_notifications-0.0.6.tar.gz", hash = "sha256:b19acac4bb41d76e59e03394b5dd165f4761c86fa327f56307a35cbee3b60158", size = 2517, upload-time = "2026-04-24T08:43:57.603Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7c/f0/21f81892e4ed10f4ec3ef2e7cf8635fb76e7c0907c55d0da66be50094760/farama_notifications-0.0.6-py3-none-any.whl", hash = "sha256:f84839188efa1ce5bb361c2a84881b2dc2c0d0d7fb661ff00421820170930935", size = 2897, upload-time = "2026-04-24T08:43:56.785Z" }, +] + +[[package]] +name = "flatbuffers" +version = "25.12.19" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/2d/d2a548598be01649e2d46231d151a6c56d10b964d94043a335ae56ea2d92/flatbuffers-25.12.19-py2.py3-none-any.whl", hash = "sha256:7634f50c427838bb021c2d66a3d1168e9d199b0607e6329399f04846d42e20b4", size = 26661, upload-time = "2025-12-19T23:16:13.622Z" }, +] + +[[package]] +name = "flax" +version = "0.10.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jax" }, + { name = "msgpack" }, + { name = "numpy" }, + { name = "optax" }, + { name = "orbax-checkpoint" }, + { name = "pyyaml" }, + { name = "rich" }, + { name = "tensorstore" }, + { name = "treescope" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/45/5d/6fa18b5dcb737f0c454ad3c116c4faebbb51f164d81109255391a0462c5e/flax-0.10.4.tar.gz", hash = "sha256:57ae44d3f111fc85cff9049adb9684ce8ebd44e87bd8ca776ed52422c2d85021", size = 5179580, upload-time = "2025-02-27T00:11:00.913Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/02/d3/c2de519b6b1f0292727406c5c12964fbe0f46bdbdbf92857716e358f1df3/flax-0.10.4-py3-none-any.whl", hash = "sha256:8cc83d91654ff943909730e02e858b4cd4577531373f83abe6597c58c581032d", size = 451817, upload-time = "2025-02-27T00:10:58.636Z" }, +] + +[[package]] +name = "fsspec" +version = "2026.9.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/77/cd/9be253869fc42e764de7f3dedd6969af7d44ff9c3375214a3442a6f3fc08/fsspec-2026.9.0.tar.gz", hash = "sha256:0f08147951c8cb31d844c3547d631053b127863b60be04cf06e121333ee0e2fe", size = 333545, upload-time = "2026-09-18T17:50:42.825Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6c/c0/a98505f18594f1bce828bb159cec0fcf9860562f1a2c85913409fc8f3d9e/fsspec-2026.9.0-py3-none-any.whl", hash = "sha256:8dd6e646e99ea382bd85f97a45e6b526a442d79423a7dc673f1e2756d05fcb5f", size = 221738, upload-time = "2026-09-18T17:50:41.341Z" }, +] + +[[package]] +name = "gast" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/91/f6/e73969782a2ecec280f8a176f2476149dd9dba69d5f8779ec6108a7721e6/gast-0.7.0.tar.gz", hash = "sha256:0bb14cd1b806722e91ddbab6fb86bba148c22b40e7ff11e248974e04c8adfdae", size = 33630, upload-time = "2025-11-29T15:30:05.266Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1d/33/f1c6a276de27b7d7339a34749cc33fa87f077f921969c47185d34a887ae2/gast-0.7.0-py3-none-any.whl", hash = "sha256:99cbf1365633a74099f69c59bd650476b96baa5ef196fec88032b00b31ba36f7", size = 22966, upload-time = "2025-11-29T15:30:03.983Z" }, +] + +[[package]] +name = "glfw" +version = "2.11.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b2/89/fa9cb2611314b84ffa68d76f154e1e508971d6e402b14bd2152684e35e5d/glfw-2.11.0.tar.gz", hash = "sha256:4bd3a462ed02ae25455dceaeb76d154faf08b48962c0cb2cf685419a34a71305", size = 36404, upload-time = "2026-10-16T20:18:53.75Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0c/65/c728be098d0a2fb1c81070386384640e577ec2db096e885d14ba8c3a0ffd/glfw-2.11.0-py2.py3-none-macosx_10_6_intel.whl", hash = "sha256:4f40b141657b014148893f3b8950906389f9d4d3b2bced41fbefcdcda776e9ad", size = 110341, upload-time = "2026-10-16T08:46:24.538Z" }, + { url = "https://files.pythonhosted.org/packages/ac/2c/dc36e841a6a0db69c19db71a8fb4bbb211924bcf587735caea8afdd230ed/glfw-2.11.0-py2.py3-none-macosx_11_0_arm64.whl", hash = "sha256:7405d3e3041add442e40a15ef5a880178ff90f486d576a82e112240c66f59555", size = 107194, upload-time = "2026-10-16T08:46:25.797Z" }, + { url = "https://files.pythonhosted.org/packages/b5/ab/59bfd8fb895c28f7a7c3504307d73b3eb18b85ec53a7025f9463400c4699/glfw-2.11.0-py2.py3-none-manylinux2014_aarch64.whl", hash = "sha256:de02fcb11126242a27f5f711b31ab83ae4dfe3567caa32dc2380e1f6050c59f1", size = 235053, upload-time = "2026-10-16T08:46:27.1Z" }, + { url = "https://files.pythonhosted.org/packages/b2/09/e3a95b97b4b174686e1b0340f9ee5b283ed65bb7cc09089fb228d36c351b/glfw-2.11.0-py2.py3-none-manylinux2014_x86_64.whl", hash = "sha256:a08e1135d02c9408d7622db3d391e3e33e4187754ec45633e9a482c775f4846c", size = 247001, upload-time = "2026-10-16T08:46:28.442Z" }, + { url = "https://files.pythonhosted.org/packages/ca/36/ca10753be547d78250aed12532ec8839f19bbeb5adac11ee8153ef59ce4f/glfw-2.11.0-py2.py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:e2f6be6bb463f98b7468386ec2e5328ee5975800580c28816ec549b4c62825c4", size = 236068, upload-time = "2026-10-16T08:46:29.587Z" }, + { url = "https://files.pythonhosted.org/packages/cc/b2/ee7afd7e84bc1a8a490710c3c0a7ad472b57f17909549dc7400d17f19c8b/glfw-2.11.0-py2.py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:76484fae4a4ef0f778a05e02422ecbe6960ae23039087fe2e0ad9742ab8c7bbe", size = 248539, upload-time = "2026-10-16T08:46:30.8Z" }, + { url = "https://files.pythonhosted.org/packages/be/b2/feb54341a0ad2c3a57ba462bcdef5543569a7d2a1b8751a9c46050c0e255/glfw-2.11.0-py2.py3-none-win32.whl", hash = "sha256:9aac61d97f71088b16322a4a75cdf72ac44ea3f1c4e7e326e5dd7c427e020903", size = 557704, upload-time = "2026-10-16T08:46:32.228Z" }, + { url = "https://files.pythonhosted.org/packages/78/d6/b0d9d83177397c018efa3579c54990f8d9b3af17fdcf97b0e549609a0571/glfw-2.11.0-py2.py3-none-win_amd64.whl", hash = "sha256:bd39c1adb57b7273813e33b9768bed88d1ea656abfc914b4a6424ec0c5d62674", size = 564482, upload-time = "2026-10-16T08:46:34.015Z" }, + { url = "https://files.pythonhosted.org/packages/03/29/5c890d0da8258c64ec7a2fb76eb4c9aebe89a38abf1aecd8683b0b4c40e6/glfw-2.11.0-py3-none-win_arm64.whl", hash = "sha256:b986dfb457d27852593686fd934c8808986ae2e4cae44ec1e8e89b1e46fc36ee", size = 174351, upload-time = "2026-10-16T08:46:35.244Z" }, +] + +[[package]] +name = "google-pasta" +version = "0.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "six" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/35/4a/0bd53b36ff0323d10d5f24ebd67af2de10a1117f5cf4d7add90df92756f1/google-pasta-0.2.0.tar.gz", hash = "sha256:c9f2c8dfc8f96d0d5808299920721be30c9eec37f2389f28904f454565c8a16e", size = 40430, upload-time = "2020-03-13T18:57:50.34Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a3/de/c648ef6835192e6e2cc03f40b19eeda4382c49b5bafb43d88b931c4c74ac/google_pasta-0.2.0-py3-none-any.whl", hash = "sha256:b32482794a366b5366a32c92a9a9201b107821889935a02b3e51f6b432ea84ed", size = 57471, upload-time = "2020-03-13T18:57:48.872Z" }, +] + +[[package]] +name = "grpcio" +version = "1.85.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5b/34/60907c03ef43039c1e91ed5bdfae41ada5648264f013cabc3f186600e571/grpcio-1.85.0.tar.gz", hash = "sha256:8aa2647afcbe5a87783e9626a2ece6277db33c6604c21690d19b5bd15609e36d", size = 14084753, upload-time = "2026-10-26T06:20:17.622Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d4/67/75531891c7f21dda119186151a5d24b257d094257e6dfdf6f6fa30d44feb/grpcio-1.85.0-cp312-cp312-linux_armv7l.whl", hash = "sha256:c02d956990c55730efcca487a4eac64f8e0fcb3af7e53c9e89cd02a985669709", size = 6795713, upload-time = "2026-10-26T06:18:46.285Z" }, + { url = "https://files.pythonhosted.org/packages/64/64/6dbc16efc2ff11d4b0b9c5701ab55fb03f9c45e61e2ae6896fbfb6769848/grpcio-1.85.0-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:b10f0f0a34ce6b122dbf5dcf8226cbe41509c602af926b1e55113ad9e49cc5f7", size = 12631580, upload-time = "2026-10-26T06:18:48.621Z" }, + { url = "https://files.pythonhosted.org/packages/d9/ff/be3621580c361eae4a4669580e9511f11e09bd25d2e9a034076ec1dd1525/grpcio-1.85.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b1ea46f63a7a9e19ce5cc09be0a6f9b347e0dcc14c4ab47b48488f6c093b8171", size = 7169722, upload-time = "2026-10-26T06:18:51.828Z" }, + { url = "https://files.pythonhosted.org/packages/b1/75/d26f7dc4d4d562aeeb82ff66107b5fee3fe69462457813e259b9d8dbdfcc/grpcio-1.85.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:e2f1c3089ed6ae75e24529a498a6b65941fdb6856c3416d62c88da944947fc88", size = 8179643, upload-time = "2026-10-26T06:18:53.878Z" }, + { url = "https://files.pythonhosted.org/packages/10/1e/4be389b352c79fc6a457ca8bfda6b0b5e8a1e0ed8eaca21f7a2d1a40a7c5/grpcio-1.85.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4ddfe8b667889b9be5990ae674b13f0db0a9eb61972306038d0140a1da0f0b87", size = 7361507, upload-time = "2026-10-26T06:18:56.002Z" }, + { url = "https://files.pythonhosted.org/packages/8b/ea/1dcc9f8e3439bacc95ef2198fafad97b2024f017af38353a650864780e66/grpcio-1.85.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:17641f3a6ff3512c34e5c367805d2632a43a773d7990f1c4c55243df6bee8fbb", size = 7928329, upload-time = "2026-10-26T06:18:58.189Z" }, + { url = "https://files.pythonhosted.org/packages/90/bd/a37f748e3c76e1b9314a2c28a89c80689ba18132e94c5d403ebc7fabdb9f/grpcio-1.85.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:1a94f32df519cb7efe3822f7f3c2daee0b0e2f01a281e1db5dbfb0efd1635b6e", size = 9180575, upload-time = "2026-10-26T06:19:00.234Z" }, + { url = "https://files.pythonhosted.org/packages/43/25/36e61d090c94108307a9eea2df00e7e2e23e3dac8b816f8ec2ce371925f0/grpcio-1.85.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:8572f9d7282cded6b018fb0147d09bbe7e59177ba6c452f1903682a80dc0857e", size = 8318253, upload-time = "2026-10-26T06:19:02.789Z" }, + { url = "https://files.pythonhosted.org/packages/c6/11/51c506bf6d104918d30cbc85b3d952e4378d038185dcf0322a93eaef7f16/grpcio-1.85.0-cp312-cp312-win32.whl", hash = "sha256:25007d0e4657dcdafa9a80677895b71de8c1d6335afb7fa6a4b35330e2d3127a", size = 4767975, upload-time = "2026-10-26T06:19:04.832Z" }, + { url = "https://files.pythonhosted.org/packages/a9/48/6d62ac3a289656c83fdc10fb895ecb9a59f8ee0318d654fb4cd20ed7293c/grpcio-1.85.0-cp312-cp312-win_amd64.whl", hash = "sha256:d6c43a06a059a95f03013aa5fce9106a6a2b6b78cfff4003a53e0eb410c0e1d7", size = 5607032, upload-time = "2026-10-26T06:19:06.876Z" }, +] + +[[package]] +name = "gymnasium" +version = "1.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cloudpickle" }, + { name = "farama-notifications" }, + { name = "numpy" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fb/3a/6713c8a92c259fd7619dc93ff19e6e435463f20d3d7b180b14ce99074c10/gymnasium-1.4.0.tar.gz", hash = "sha256:9754f630a32abfdbb76386abe1bc1e706c982db2d62dfa119c9952eb2de7697d", size = 350436, upload-time = "2026-10-05T11:02:16.217Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/38/4e/6a51be94ae8e9dc8eaddd6ea74a6dde754448c4c746b51a232681f81e41c/gymnasium-1.4.0-py3-none-any.whl", hash = "sha256:1cb947c59e7c72d8eabb2c2274c00cd20948e69f8452a9ef4092223be24fdf0e", size = 476306, upload-time = "2026-10-05T11:02:14.783Z" }, +] + +[package.optional-dependencies] +mujoco = [ + { name = "imageio" }, + { name = "mujoco" }, + { name = "packaging" }, +] + +[[package]] +name = "h5py" +version = "3.16.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/db/33/acd0ce6863b6c0d7735007df01815403f5589a21ff8c2e1ee2587a38f548/h5py-3.16.0.tar.gz", hash = "sha256:a0dbaad796840ccaa67a4c144a0d0c8080073c34c76d5a6941d6818678ef2738", size = 446526, upload-time = "2026-03-06T13:49:08.07Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/c0/5d4119dba94093bbafede500d3defd2f5eab7897732998c04b54021e530b/h5py-3.16.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:c5313566f4643121a78503a473f0fb1e6dcc541d5115c44f05e037609c565c4d", size = 3685604, upload-time = "2026-03-06T13:48:04.198Z" }, + { url = "https://files.pythonhosted.org/packages/b0/42/c84efcc1d4caebafb1ecd8be4643f39c85c47a80fe254d92b8b43b1eadaf/h5py-3.16.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:42b012933a83e1a558c673176676a10ce2fd3759976a0fedee1e672d1e04fc9d", size = 3061940, upload-time = "2026-03-06T13:48:05.783Z" }, + { url = "https://files.pythonhosted.org/packages/89/84/06281c82d4d1686fde1ac6b0f307c50918f1c0151062445ab3b6fa5a921d/h5py-3.16.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:ff24039e2573297787c3063df64b60aab0591980ac898329a08b0320e0cf2527", size = 5198852, upload-time = "2026-03-06T13:48:07.482Z" }, + { url = "https://files.pythonhosted.org/packages/9e/e9/1a19e42cd43cc1365e127db6aae85e1c671da1d9a5d746f4d34a50edb577/h5py-3.16.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:dfc21898ff025f1e8e67e194965a95a8d4754f452f83454538f98f8a3fcb207e", size = 5405250, upload-time = "2026-03-06T13:48:09.628Z" }, + { url = "https://files.pythonhosted.org/packages/b7/8e/9790c1655eabeb85b92b1ecab7d7e62a2069e53baefd58c98f0909c7a948/h5py-3.16.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:698dd69291272642ffda44a0ecd6cd3bda5faf9621452d255f57ce91487b9794", size = 5190108, upload-time = "2026-03-06T13:48:11.26Z" }, + { url = "https://files.pythonhosted.org/packages/51/d7/ab693274f1bd7e8c5f9fdd6c7003a88d59bedeaf8752716a55f532924fbb/h5py-3.16.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:2b2c02b0a160faed5fb33f1ba8a264a37ee240b22e049ecc827345d0d9043074", size = 5419216, upload-time = "2026-03-06T13:48:13.322Z" }, + { url = "https://files.pythonhosted.org/packages/03/c1/0976b235cf29ead553e22f2fb6385a8252b533715e00d0ae52ed7b900582/h5py-3.16.0-cp312-cp312-win_amd64.whl", hash = "sha256:96b422019a1c8975c2d5dadcf61d4ba6f01c31f92bbde6e4649607885fe502d6", size = 3182868, upload-time = "2026-03-06T13:48:15.759Z" }, + { url = "https://files.pythonhosted.org/packages/14/d9/866b7e570b39070f92d47b0ff1800f0f8239b6f9e45f02363d7112336c1f/h5py-3.16.0-cp312-cp312-win_arm64.whl", hash = "sha256:39c2838fb1e8d97bcf1755e60ad1f3dd76a7b2a475928dc321672752678b96db", size = 2653286, upload-time = "2026-03-06T13:48:17.279Z" }, +] + +[[package]] +name = "humanize" +version = "4.16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0a/ea/13a1ef3c12d12662905801495283530251918b70d62d368f1d2e0272c70d/humanize-4.16.0.tar.gz", hash = "sha256:7dc2244a2f84a4bfb1d36c37bac80cd78e35cdc5c119206d87b018e1445f3a3f", size = 89515, upload-time = "2026-06-30T16:17:29.859Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b0/aa/0b7365d30fed43e7a3449aba1fe20a0a7174d9cf13e282af4e69ac825441/humanize-4.16.0-py3-none-any.whl", hash = "sha256:353eb2f34c09d098b2880eee8bef21832eae6d174f48c5762fff7e5fcb74d01d", size = 137209, upload-time = "2026-06-30T16:17:28.36Z" }, +] + +[[package]] +name = "idna" +version = "3.20" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f5/08/8eea9d4b8302028f3abb2c0813953f7aec26d33b7a8960ed760e65ff29fa/idna-3.20.tar.gz", hash = "sha256:a7db850025b95ded1eae8a46181a1a6c56c92c96f0e2b005d9ff8dc0210cab44", size = 216463, upload-time = "2026-09-17T14:11:04.752Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/58/a2/bb081bab032533a855d44de1d56f8e8426114ff1ba5d1f07a438a0a654f8/idna-3.20-py3-none-any.whl", hash = "sha256:ab7ae7122974553370f0bdb919e1a960b2cd1bc1ef0276416d896db81c14582c", size = 69583, upload-time = "2026-09-17T14:11:03.168Z" }, +] + +[[package]] +name = "imageio" +version = "2.38.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "pillow" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f3/cd/69e4ac55b6dafdd2b5f32075236841a3945dea7323d0232d80f28c37cea8/imageio-2.38.1.tar.gz", hash = "sha256:6769f1f01c4dd46448307863a787c9a22fa4dbe11c0c88f525c6e410fdfd7983", size = 403667, upload-time = "2026-10-08T14:35:32.961Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7d/8a/3b7f62b9df56460959f1b9bef094248d2a4d7f393c5b0d2da8efcdf0e40e/imageio-2.38.1-py3-none-any.whl", hash = "sha256:36d23eb7423d2fb63f637098758edb3d2df3687f7125e0a5ce8596572022dcf9", size = 328389, upload-time = "2026-10-08T14:35:31.075Z" }, +] + +[[package]] +name = "iniconfig" +version = "2.3.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/01/e1/2069291243c926a2ff1cd706c7f3eeb9b62144bf60f77c9fb9ff2fb26bd3/iniconfig-2.3.1.tar.gz", hash = "sha256:67f4b9c50da0dedf52af349e7749a80a9057a5031199791b906c3bb3ae878960", size = 21209, upload-time = "2026-10-06T22:48:38.076Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/43/4ca9e49d27a1fcf6bece6f6aec0ea46bb9112489b93d4b688fb415457bdb/iniconfig-2.3.1-py3-none-any.whl", hash = "sha256:9121e2c1fdb355232495be3194c8dfe87ccc2d5dee45947b78e68f499790d7a7", size = 7552, upload-time = "2026-10-06T22:48:36.959Z" }, +] + +[[package]] +name = "jax" +version = "0.4.35" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jaxlib" }, + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "opt-einsum" }, + { name = "scipy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e3/34/21da583b9596e72bb8e95b6197dee0a44b96b9ea2c147fccabd43ca5515b/jax-0.4.35.tar.gz", hash = "sha256:c0c986993026b10bf6f607fecb7417377460254640766ce40f1fef3fd139c12e", size = 1861189, upload-time = "2024-10-22T20:56:36.379Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/20/6c57c50c0ccc645fea1895950f1e5cd02f961ee44b3ffe83617fa46b0c1d/jax-0.4.35-py3-none-any.whl", hash = "sha256:fa99e909a31424abfec750019a6dd36f6acc18a6e7d40e2c0086b932cc351325", size = 2158621, upload-time = "2024-10-22T20:56:34.071Z" }, +] + +[package.optional-dependencies] +cuda12 = [ + { name = "jax-cuda12-plugin", extra = ["with-cuda"] }, + { name = "jaxlib" }, +] + +[[package]] +name = "jax-cuda12-pjrt" +version = "0.4.35" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8f/ba/12d0f831331e0eb7dafc49eb1bd0a905098aca917456d6ccf57efd57b908/jax_cuda12_pjrt-0.4.35-py3-none-manylinux2014_aarch64.whl", hash = "sha256:81980f1696749351364d7999ee5cb21bedd9d14daa01b40eb60d20f2210da11c", size = 84301075, upload-time = "2024-10-22T20:58:42.787Z" }, + { url = "https://files.pythonhosted.org/packages/d7/aa/f15ea857ad9bcff7a0c942dc570ca718b026cc0cc5c513525bb08cacf3c0/jax_cuda12_pjrt-0.4.35-py3-none-manylinux2014_x86_64.whl", hash = "sha256:0ffe7e1ba65659bd5738c2cc5addaf0a56205d2188eec5da194b63c068e1fdd2", size = 100760554, upload-time = "2024-10-22T20:58:49.24Z" }, +] + +[[package]] +name = "jax-cuda12-plugin" +version = "0.4.35" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jax-cuda12-pjrt" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/0f/6b/2b08b0a33dbbd410ed6746d2f12c67b4b7a1318590cf05fa3ca7dbff4c46/jax_cuda12_plugin-0.4.35-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:767e53cd0ab0f543f7ac64c74e326b4e7145ca6424d0e7128373e5777012e6d9", size = 15541645, upload-time = "2024-10-22T22:36:38.277Z" }, + { url = "https://files.pythonhosted.org/packages/23/a1/781f98da41b9abf63a1468c5babbafa7b1bdb766c081386644a4034fbdef/jax_cuda12_plugin-0.4.35-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:80aa5e628e8e9dd4f89b767f73900152033c3bf6d5f49dc2081e93f9bac1c9dc", size = 15516745, upload-time = "2024-10-22T22:36:40.825Z" }, +] + +[package.optional-dependencies] +with-cuda = [ + { name = "nvidia-cublas-cu12" }, + { name = "nvidia-cuda-cupti-cu12" }, + { name = "nvidia-cuda-nvcc-cu12" }, + { name = "nvidia-cuda-runtime-cu12" }, + { name = "nvidia-cudnn-cu12" }, + { name = "nvidia-cufft-cu12" }, + { name = "nvidia-cusolver-cu12" }, + { name = "nvidia-cusparse-cu12" }, + { name = "nvidia-nccl-cu12" }, + { name = "nvidia-nvjitlink-cu12" }, +] + +[[package]] +name = "jaxlib" +version = "0.4.34" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "scipy" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/42/5038983664494dfb50f8669a662d965d7ea62f9250e40d8cd36dcf9ac3dd/jaxlib-0.4.34-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:c7b3e724a30426a856070aba0192b5d199e95b4411070e7ad96ad8b196877b10", size = 87473956, upload-time = "2024-10-04T14:38:24.126Z" }, + { url = "https://files.pythonhosted.org/packages/87/2e/8a75d3107c019c370c50c01acc205da33f9d6fba830950401a772a8e9f6d/jaxlib-0.4.34-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:096f0ca309d41fa692a9d1f2f9baab1c5c8ca0749876ebb3f748e738a27c7ff4", size = 67650276, upload-time = "2024-10-04T14:38:29.587Z" }, + { url = "https://files.pythonhosted.org/packages/af/09/cceae2d251a506b4297679d10ee9f5e905a6b992b0687d553c9470ffd1db/jaxlib-0.4.34-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:1a30771d85fa77f9ab8f18e63240f455ab3a3f87660ed7b8d5eea6ceecbe5c1e", size = 69431284, upload-time = "2024-10-04T14:38:34.902Z" }, + { url = "https://files.pythonhosted.org/packages/e7/0d/4faf839e3c8ce2a5b615df64427be3e870899c72c0ebfb5859348150aba1/jaxlib-0.4.34-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:48272e9034ff868d4328cf0055a07882fd2be93f59dfb6283af7de491f9d1290", size = 86151183, upload-time = "2024-10-04T14:38:40.416Z" }, + { url = "https://files.pythonhosted.org/packages/a4/bc/a38f99071fca6cc31ae949e508a23b0de5de559da594443bb625a1adb8f3/jaxlib-0.4.34-cp312-cp312-win_amd64.whl", hash = "sha256:901cb4040ed24eae40071d8114ea8d10dff436277fa74a1a5b9e7206f641151c", size = 55278745, upload-time = "2024-10-04T14:38:45.918Z" }, +] + +[[package]] +name = "keras" +version = "3.15.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "h5py" }, + { name = "ml-dtypes" }, + { name = "namex" }, + { name = "numpy" }, + { name = "optree" }, + { name = "packaging" }, + { name = "rich" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ec/a6/4cebb4196f4e6284b35e68dc17dda9a4111a2e2bf21db8211de12881ef89/keras-3.15.1.tar.gz", hash = "sha256:92ae7c1dd7f61041953dc2d42253181ee099086f0b58ae36fcf98147a6f08a29", size = 1919818, upload-time = "2026-07-29T18:20:13.909Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/b3/c9b848bbdba18e765a8051917c0cc82585b64278ce87895f2e521a27438f/keras-3.15.1-py3-none-any.whl", hash = "sha256:836460e480930acbd19bb7a17e62f9ecad40e8a9af9a651fc8d0586a96d27e20", size = 2398595, upload-time = "2026-07-29T18:20:12.345Z" }, +] + +[[package]] +name = "libclang" +version = "18.1.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6e/5c/ca35e19a4f142adffa27e3d652196b7362fa612243e2b916845d801454fc/libclang-18.1.1.tar.gz", hash = "sha256:a1214966d08d73d971287fc3ead8dfaf82eb07fb197680d8b3859dbbbbf78250", size = 39612, upload-time = "2024-03-17T16:04:37.434Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4b/49/f5e3e7e1419872b69f6f5e82ba56e33955a74bd537d8a1f5f1eff2f3668a/libclang-18.1.1-1-py2.py3-none-macosx_11_0_arm64.whl", hash = "sha256:0b2e143f0fac830156feb56f9231ff8338c20aecfe72b4ffe96f19e5a1dbb69a", size = 25836045, upload-time = "2024-06-30T17:40:31.646Z" }, + { url = "https://files.pythonhosted.org/packages/e2/e5/fc61bbded91a8830ccce94c5294ecd6e88e496cc85f6704bf350c0634b70/libclang-18.1.1-py2.py3-none-macosx_10_9_x86_64.whl", hash = "sha256:6f14c3f194704e5d09769108f03185fce7acaf1d1ae4bbb2f30a72c2400cb7c5", size = 26502641, upload-time = "2024-03-18T15:52:26.722Z" }, + { url = "https://files.pythonhosted.org/packages/db/ed/1df62b44db2583375f6a8a5e2ca5432bbdc3edb477942b9b7c848c720055/libclang-18.1.1-py2.py3-none-macosx_11_0_arm64.whl", hash = "sha256:83ce5045d101b669ac38e6da8e58765f12da2d3aafb3b9b98d88b286a60964d8", size = 26420207, upload-time = "2024-03-17T15:00:26.63Z" }, + { url = "https://files.pythonhosted.org/packages/1d/fc/716c1e62e512ef1c160e7984a73a5fc7df45166f2ff3f254e71c58076f7c/libclang-18.1.1-py2.py3-none-manylinux2010_x86_64.whl", hash = "sha256:c533091d8a3bbf7460a00cb6c1a71da93bffe148f172c7d03b1c31fbf8aa2a0b", size = 24515943, upload-time = "2024-03-17T16:03:45.942Z" }, + { url = "https://files.pythonhosted.org/packages/3c/3d/f0ac1150280d8d20d059608cf2d5ff61b7c3b7f7bcf9c0f425ab92df769a/libclang-18.1.1-py2.py3-none-manylinux2014_aarch64.whl", hash = "sha256:54dda940a4a0491a9d1532bf071ea3ef26e6dbaf03b5000ed94dd7174e8f9592", size = 23784972, upload-time = "2024-03-17T16:12:47.677Z" }, + { url = "https://files.pythonhosted.org/packages/fe/2f/d920822c2b1ce9326a4c78c0c2b4aa3fde610c7ee9f631b600acb5376c26/libclang-18.1.1-py2.py3-none-manylinux2014_armv7l.whl", hash = "sha256:cf4a99b05376513717ab5d82a0db832c56ccea4fd61a69dbb7bccf2dfb207dbe", size = 20259606, upload-time = "2024-03-17T16:17:42.437Z" }, + { url = "https://files.pythonhosted.org/packages/2d/c2/de1db8c6d413597076a4259cea409b83459b2db997c003578affdd32bf66/libclang-18.1.1-py2.py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:69f8eb8f65c279e765ffd28aaa7e9e364c776c17618af8bff22a8df58677ff4f", size = 24921494, upload-time = "2024-03-17T16:14:20.132Z" }, + { url = "https://files.pythonhosted.org/packages/0b/2d/3f480b1e1d31eb3d6de5e3ef641954e5c67430d5ac93b7fa7e07589576c7/libclang-18.1.1-py2.py3-none-win_amd64.whl", hash = "sha256:4dd2d3b82fab35e2bf9ca717d7b63ac990a3519c7e312f19fa8e86dcc712f7fb", size = 26415083, upload-time = "2024-03-17T16:42:21.703Z" }, + { url = "https://files.pythonhosted.org/packages/71/cf/e01dc4cc79779cd82d77888a88ae2fa424d93b445ad4f6c02bfc18335b70/libclang-18.1.1-py2.py3-none-win_arm64.whl", hash = "sha256:3f0e1f49f04d3cd198985fea0511576b0aee16f9ff0e0f0cad7f9c57ec3c20e8", size = 22361112, upload-time = "2024-03-17T16:42:59.565Z" }, +] + +[[package]] +name = "markdown" +version = "3.11.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f4/32/0b4d0297064244f04af4331b65de16237cc3cce07474171ee14222a28981/markdown-3.11.2.tar.gz", hash = "sha256:87bf5d90f2d2882e2c02fcda182c2800456b985bff38913867c49706481514b0", size = 491445, upload-time = "2026-10-19T19:41:42.687Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/66/bcab3b677b71ad286b4e607e3427499537c76e13ad2e22584ba7a7d9882c/markdown-3.11.2-py3-none-any.whl", hash = "sha256:f80390747a2a0ff2ba54da074ce4f862f72a6685f608cd2637b646a6d83993f5", size = 114266, upload-time = "2026-10-19T19:41:41.244Z" }, +] + +[[package]] +name = "markdown-it-py" +version = "4.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mdurl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/ff/7841249c247aa650a76b9ee4bbaeae59370dc8bfd2f6c01f3630c35eb134/markdown_it_py-4.2.0.tar.gz", hash = "sha256:04a21681d6fbb623de53f6f364d352309d4094dd4194040a10fd51833e418d49", size = 82454, upload-time = "2026-05-07T12:08:28.36Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/81/4da04ced5a082363ecfa159c010d200ecbd959ae410c10c0264a38cac0f5/markdown_it_py-4.2.0-py3-none-any.whl", hash = "sha256:9f7ebbcd14fe59494226453aed97c1070d83f8d24b6fc3a3bcf9a38092641c4a", size = 91687, upload-time = "2026-05-07T12:08:27.182Z" }, +] + +[[package]] +name = "markupsafe" +version = "3.0.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/38/9b/e422a865e1d5d57d0e509b4e0bf1c1a70a7f6382c29a5aa428df994c8bc8/markupsafe-3.0.4.tar.gz", hash = "sha256:2e9ad7dd851bf45fab9f75cbff4cb493fee9979e8d8c7c9c3ee119022518edd6", size = 153777, upload-time = "2026-10-02T23:07:22.29Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/81/09/4c59d56b8461ae8eb0d8ba34bb25b7e618547044679d58a82ef9b2479fc1/markupsafe-3.0.4-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:61631e08084be9e21a8967ec3139c7616ed7c5e9368e05c86d1b39562c8a57b6", size = 11658, upload-time = "2026-10-02T23:04:51.876Z" }, + { url = "https://files.pythonhosted.org/packages/a2/f0/d6613774d86fbf6d145751d43c59875e47a6f9f17daee0aef173bd36d90e/markupsafe-3.0.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0930db9bdc62d22944e10b066448bb65dc9abe9112880c7cab8da54db4284d5f", size = 12050, upload-time = "2026-10-02T23:04:52.931Z" }, + { url = "https://files.pythonhosted.org/packages/0d/f2/8f18e0b806eb13c1f8d07d917a720831ead54253a6dec011fbc78098a6f8/markupsafe-3.0.4-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6a45c3d514f2436064db00d7fc8778d888f0236ebfed649b53d13a59e69ad51b", size = 24363, upload-time = "2026-10-02T23:04:53.895Z" }, + { url = "https://files.pythonhosted.org/packages/60/ce/fa07dbe8a5675558fa36dea033e19995bc783de2dec5f540ccb9030b06aa/markupsafe-3.0.4-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:1e1451fab512d1bcc3dc26988ec1edb0b82c2db909132872cd9356070a6b63df", size = 28525, upload-time = "2026-10-02T23:04:54.905Z" }, + { url = "https://files.pythonhosted.org/packages/85/40/be87c01f3868ec217f8a2015089d71c22c8c5a75324822e5ed1cdd87210d/markupsafe-3.0.4-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:bd3ce56ae2cbae3ba82b683bc425cd7e48d2ed8b10f3e818186b6f5646d9271c", size = 24733, upload-time = "2026-10-02T23:04:56.229Z" }, + { url = "https://files.pythonhosted.org/packages/4f/a7/aeedb5140afa41fc74c225e9184ab96723a6e873b6ee1c9fede7283456d8/markupsafe-3.0.4-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8e124f974786f831d6043728e38296969d3579db8896fe004682f5758e613581", size = 22985, upload-time = "2026-10-02T23:04:57.521Z" }, + { url = "https://files.pythonhosted.org/packages/c3/fc/e91352bb08c6a59da3ef0909d457bf95a5f5908fbf151b30a06d9dbcfbb4/markupsafe-3.0.4-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:c02e8f18bdedba082cef725942ac823b9b60656db07f7e265cb31618dfd00d77", size = 22001, upload-time = "2026-10-02T23:04:58.597Z" }, + { url = "https://files.pythonhosted.org/packages/5d/f8/bffee5e7d2a3deb59748a797650a48af7e672025cf641a79344a771ad106/markupsafe-3.0.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:9f098115c247e11d138ab83a28fa0323c77015007ea2df73ba5fd714dfefd67c", size = 23793, upload-time = "2026-10-02T23:04:59.686Z" }, + { url = "https://files.pythonhosted.org/packages/ed/59/b853d6628ecb4d658e1d637224846d5e9bb4adf4f8df97f3be9f29dce2ec/markupsafe-3.0.4-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:d5f93ebbeb8032d47e349328ec8662d973d9b05a70b3c35df1f91fe419b84749", size = 22640, upload-time = "2026-10-02T23:05:00.768Z" }, + { url = "https://files.pythonhosted.org/packages/09/b2/1506df394f0f075797c418d0301498f49e43be194e3ffcb49e6fe6ccf022/markupsafe-3.0.4-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:64511c54db4e4987aef4c41923235927428729e8174c5dba488429be70a998ed", size = 24074, upload-time = "2026-10-02T23:05:01.813Z" }, + { url = "https://files.pythonhosted.org/packages/c7/81/5ed69cda630ac69ef60d06c09ba5a7f84ff66a2e28cf986fd5614ab3c6e6/markupsafe-3.0.4-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:e1a622f13970d81f95d0c72f9dc090dce9085fccfa4c9f2174377ee32bd15786", size = 21563, upload-time = "2026-10-02T23:05:03.239Z" }, + { url = "https://files.pythonhosted.org/packages/0c/fe/fb1e79be0fea60aa32602ebefc9c35a82bb42b4df157285ab7dfec12341a/markupsafe-3.0.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:c9a7f43c0b202b334cc9184af09bb8f21d3a209e038efaf106936fb69e6b026e", size = 23048, upload-time = "2026-10-02T23:05:04.479Z" }, + { url = "https://files.pythonhosted.org/packages/c8/52/7632a53360671a9b750cdbabaf9cdd89f18b42248b8e4cb42c0b0296e459/markupsafe-3.0.4-cp312-cp312-win32.whl", hash = "sha256:f0ec3b750b59375eab5b0fb2b9254810c00a3375be6d789899f1055a1d556237", size = 14108, upload-time = "2026-10-02T23:05:05.513Z" }, + { url = "https://files.pythonhosted.org/packages/3f/bf/62495e180b7000aaf30000fff849e933f74264638057176cf46852500adc/markupsafe-3.0.4-cp312-cp312-win_amd64.whl", hash = "sha256:11935df9bf455ed0c04eb87bcd720f02b1fe5e02128a9430f23aed6f93336fc7", size = 14303, upload-time = "2026-10-02T23:05:06.538Z" }, + { url = "https://files.pythonhosted.org/packages/c5/8e/4c24208776a65878d656996945aacfbfe010d3720d1a98fc0eb8491fc03b/markupsafe-3.0.4-cp312-cp312-win_arm64.whl", hash = "sha256:a4bbd2d87dd233b9fc5812160c3d0ffbe42edc22a26ce0469f58479ede633fe9", size = 14174, upload-time = "2026-10-02T23:05:07.617Z" }, +] + +[[package]] +name = "mdurl" +version = "0.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, +] + +[[package]] +name = "ml-dtypes" +version = "0.4.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fd/15/76f86faa0902836cc133939732f7611ace68cf54148487a99c539c272dc8/ml_dtypes-0.4.1.tar.gz", hash = "sha256:fad5f2de464fd09127e49b7fd1252b9006fb43d2edc1ff112d390c324af5ca7a", size = 692594, upload-time = "2024-09-13T19:07:11.624Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ba/1a/99e924f12e4b62139fbac87419698c65f956d58de0dbfa7c028fa5b096aa/ml_dtypes-0.4.1-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:827d3ca2097085cf0355f8fdf092b888890bb1b1455f52801a2d7756f056f54b", size = 405077, upload-time = "2024-09-13T19:06:57.538Z" }, + { url = "https://files.pythonhosted.org/packages/8f/8c/7b610bd500617854c8cc6ed7c8cfb9d48d6a5c21a1437a36a4b9bc8a3598/ml_dtypes-0.4.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:772426b08a6172a891274d581ce58ea2789cc8abc1c002a27223f314aaf894e7", size = 2181554, upload-time = "2024-09-13T19:06:59.196Z" }, + { url = "https://files.pythonhosted.org/packages/c7/c6/f89620cecc0581dc1839e218c4315171312e46c62a62da6ace204bda91c0/ml_dtypes-0.4.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:126e7d679b8676d1a958f2651949fbfa182832c3cd08020d8facd94e4114f3e9", size = 2160488, upload-time = "2024-09-13T19:07:03.131Z" }, + { url = "https://files.pythonhosted.org/packages/ae/11/a742d3c31b2cc8557a48efdde53427fd5f9caa2fa3c9c27d826e78a66f51/ml_dtypes-0.4.1-cp312-cp312-win_amd64.whl", hash = "sha256:df0fb650d5c582a9e72bb5bd96cfebb2cdb889d89daff621c8fbc60295eba66c", size = 127462, upload-time = "2024-09-13T19:07:04.916Z" }, +] + +[[package]] +name = "msgpack" +version = "1.2.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cc/52/1ec7ead2c7528bf3d1db73e886fa4df10f9e6f0eb7854a448b7090cfcef7/msgpack-1.2.4.tar.gz", hash = "sha256:475c44e12f2a5de766a8448783f149c9b33f3f49b88953c6cd739a77d93c732c", size = 197764, upload-time = "2026-10-26T17:55:48.593Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/27/fd/45f6cd13cefbb659a8e0128ea64d51b349571163fec3c9c20f99792dc42a/msgpack-1.2.4-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:90cba4996b18ec1201fcaabd01ec48ff4b9755416bcd52105ded2f89fe145c49", size = 91903, upload-time = "2026-10-26T17:54:03.097Z" }, + { url = "https://files.pythonhosted.org/packages/7f/89/169c2a0b1ffeafe68a4f13d8d96b6ed1cd8d40d8262b3a6690851da7eda1/msgpack-1.2.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:112933726919f4239646cca9572e2c2fa5626358c670da8440aac09403100a48", size = 90573, upload-time = "2026-10-26T17:54:04.383Z" }, + { url = "https://files.pythonhosted.org/packages/cf/8e/272a384f2a001cfced1baab0a8abdd912d6281432670f4609d1aab718d7b/msgpack-1.2.4-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9ca1e78106675f2aed07dc37b842ce19be70010ab22c5c5c04688f0840d48539", size = 461057, upload-time = "2026-10-26T17:54:05.821Z" }, + { url = "https://files.pythonhosted.org/packages/1e/d9/f59bd841fc2fd2b74dc50ce85860a6a25f58e588dd48946aeaf2862cf990/msgpack-1.2.4-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c4d5bbdbea8f5669b3a734d6c226772804fb120848f99cbe63a0ddee63b049c9", size = 475652, upload-time = "2026-10-26T17:54:07.414Z" }, + { url = "https://files.pythonhosted.org/packages/00/03/afd000e1b1e28fd535dc8c7bc21ebb88923bc1964a2429fe660fb313ace3/msgpack-1.2.4-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:0d9ff7f3aaffd84010f1c501b05406a497584f5cab2f732eb9d89bf960e12f6f", size = 427082, upload-time = "2026-10-26T17:54:09.1Z" }, + { url = "https://files.pythonhosted.org/packages/f7/7a/e2890b272883c496faad777324ddb70e9cc4a6c63731962a4a39a6a40f44/msgpack-1.2.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:aa10dc04c03408f31948ac1896a942141fd8d7a954b3612edac79f19663e4076", size = 451956, upload-time = "2026-10-26T17:54:10.476Z" }, + { url = "https://files.pythonhosted.org/packages/c2/46/0c1b240b8a580af3e72139c33b018025440a1a42c2579e82e39b1195ee36/msgpack-1.2.4-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:472a25463678f2795143a7cb42e88a1b59cc2375e9881917e69a9323bde42114", size = 423833, upload-time = "2026-10-26T17:54:12.195Z" }, + { url = "https://files.pythonhosted.org/packages/80/8e/851415b1a9264e6abf64d8a9617d856769d59e247fb9470106f133aedbae/msgpack-1.2.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:cbb8a677da3e845713925e9a203b8a120b1e9da9989836890ba4a4a187902701", size = 470951, upload-time = "2026-10-26T17:54:13.764Z" }, + { url = "https://files.pythonhosted.org/packages/c9/74/d75bf61aab9cb88b84f7d3a8eca8dd92e19a742b77ded300a1177e88d031/msgpack-1.2.4-cp312-cp312-win32.whl", hash = "sha256:2b0b7dff5beddc10f16ecea6b28c6b98014d1135d7015bd83adff5037f35848d", size = 68786, upload-time = "2026-10-26T17:54:15.049Z" }, + { url = "https://files.pythonhosted.org/packages/81/c5/dd0e78eb04f9a0773a810d25ef2cd8f16bb941487567d0c63f814881133c/msgpack-1.2.4-cp312-cp312-win_amd64.whl", hash = "sha256:6e38b123681df22373d545520f54bbeddc4706d9235a1161061f12bd6648497d", size = 77272, upload-time = "2026-10-26T17:54:16.491Z" }, + { url = "https://files.pythonhosted.org/packages/65/e7/b38f63ae82777c0077daa56c4b09af39d29dcb8927ab2ff59268094069c5/msgpack-1.2.4-cp312-cp312-win_arm64.whl", hash = "sha256:6795c588b210b2f13d62c8fcc23bf779b509b839ec5c197a7ab4f5191459e126", size = 71884, upload-time = "2026-10-26T17:54:17.995Z" }, +] + +[[package]] +name = "mujoco" +version = "3.17.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "etils", extra = ["epath"] }, + { name = "glfw" }, + { name = "numpy" }, + { name = "pyopengl" }, + { name = "websockets" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/40/10/dde0a84f2c4552729540c3e19431177e6e1b4c20dc536e19169e575d0e94/mujoco-3.17.0.tar.gz", hash = "sha256:5e49c15905e7a39eaa1efb1cae809a81c5f5c2439467f23e7e6e383624d957be", size = 7145175, upload-time = "2026-10-22T23:58:49.331Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/45/da116334f2e21765d9300f0810f714f8891cc73968115a57b61b0cecfe74/mujoco-3.17.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c677ab789299db70ca3fffe96b44a424bdad516cc87cc2a83898140923c2e91e", size = 24955426, upload-time = "2026-10-22T23:34:50.567Z" }, + { url = "https://files.pythonhosted.org/packages/a9/a5/d83393779c81fe1807442844254fe71b9182f1f2b64e865dfe8534db1620/mujoco-3.17.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:525e09a00ecd964f06582773d9eeb03bb0e5a1675edf296a6868024759b3b235", size = 25231702, upload-time = "2026-10-22T23:34:54.582Z" }, + { url = "https://files.pythonhosted.org/packages/66/33/9e329714eefc47dfbedefe6ea3809f055acbf58279799d635b1a18510924/mujoco-3.17.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:42371356efe53a5f5f73ed03fcb967bdf04048017f80e2fca01e71a360e42da2", size = 26654565, upload-time = "2026-10-22T23:34:57.762Z" }, + { url = "https://files.pythonhosted.org/packages/c1/16/609b306f340fbc58fb0a75a5aac7cb4bc8f9d60155ebebc22043a4526069/mujoco-3.17.0-cp312-cp312-win_amd64.whl", hash = "sha256:17ab707e86b0a5da13582a5335d3ea585a5ac4cce2e814e15b203530d1bf16e3", size = 16311296, upload-time = "2026-10-22T23:57:32.433Z" }, +] + +[[package]] +name = "namex" +version = "0.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0c/c0/ee95b28f029c73f8d49d8f52edaed02a1d4a9acb8b69355737fdb1faa191/namex-0.1.0.tar.gz", hash = "sha256:117f03ccd302cc48e3f5c58a296838f6b89c83455ab8683a1e85f2a430aa4306", size = 6649, upload-time = "2025-05-26T23:17:38.918Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b2/bc/465daf1de06409cdd4532082806770ee0d8d7df434da79c76564d0f69741/namex-0.1.0-py3-none-any.whl", hash = "sha256:e2012a474502f1e2251267062aae3114611f07df4224b6e06334c57b0f2ce87c", size = 5905, upload-time = "2025-05-26T23:17:37.695Z" }, +] + +[[package]] +name = "nest-asyncio" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/83/f8/51569ac65d696c8ecbee95938f89d4abf00f47d58d48f6fbabfe8f0baefe/nest_asyncio-1.6.0.tar.gz", hash = "sha256:6f172d5449aca15afd6c646851f4e31e02c598d553a667e38cafa997cfec55fe", size = 7418, upload-time = "2024-01-21T14:25:19.227Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/c4/c2971a3ba4c6103a3d10c4b0f24f461ddc027f0f09763220cf35ca1401b3/nest_asyncio-1.6.0-py3-none-any.whl", hash = "sha256:87af6efd6b5e897c81050477ef65c62e2b2f35d51703cae01aff2905b1852e1c", size = 5195, upload-time = "2024-01-21T14:25:17.223Z" }, +] + +[[package]] +name = "numpy" +version = "1.26.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/65/6e/09db70a523a96d25e115e71cc56a6f9031e7b8cd166c1ac8438307c14058/numpy-1.26.4.tar.gz", hash = "sha256:2a02aba9ed12e4ac4eb3ea9421c420301a0c6460d9830d74a9df87efa4912010", size = 15786129, upload-time = "2024-02-06T00:26:44.495Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/95/12/8f2020a8e8b8383ac0177dc9570aad031a3beb12e38847f7129bacd96228/numpy-1.26.4-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:b3ce300f3644fb06443ee2222c2201dd3a89ea6040541412b8fa189341847218", size = 20335901, upload-time = "2024-02-05T23:55:32.801Z" }, + { url = "https://files.pythonhosted.org/packages/75/5b/ca6c8bd14007e5ca171c7c03102d17b4f4e0ceb53957e8c44343a9546dcc/numpy-1.26.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:03a8c78d01d9781b28a6989f6fa1bb2c4f2d51201cf99d3dd875df6fbd96b23b", size = 13685868, upload-time = "2024-02-05T23:55:56.28Z" }, + { url = "https://files.pythonhosted.org/packages/79/f8/97f10e6755e2a7d027ca783f63044d5b1bc1ae7acb12afe6a9b4286eac17/numpy-1.26.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9fad7dcb1aac3c7f0584a5a8133e3a43eeb2fe127f47e3632d43d677c66c102b", size = 13925109, upload-time = "2024-02-05T23:56:20.368Z" }, + { url = "https://files.pythonhosted.org/packages/0f/50/de23fde84e45f5c4fda2488c759b69990fd4512387a8632860f3ac9cd225/numpy-1.26.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:675d61ffbfa78604709862923189bad94014bef562cc35cf61d3a07bba02a7ed", size = 17950613, upload-time = "2024-02-05T23:56:56.054Z" }, + { url = "https://files.pythonhosted.org/packages/4c/0c/9c603826b6465e82591e05ca230dfc13376da512b25ccd0894709b054ed0/numpy-1.26.4-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:ab47dbe5cc8210f55aa58e4805fe224dac469cde56b9f731a4c098b91917159a", size = 13572172, upload-time = "2024-02-05T23:57:21.56Z" }, + { url = "https://files.pythonhosted.org/packages/76/8c/2ba3902e1a0fc1c74962ea9bb33a534bb05984ad7ff9515bf8d07527cadd/numpy-1.26.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:1dda2e7b4ec9dd512f84935c5f126c8bd8b9f2fc001e9f54af255e8c5f16b0e0", size = 17786643, upload-time = "2024-02-05T23:57:56.585Z" }, + { url = "https://files.pythonhosted.org/packages/28/4a/46d9e65106879492374999e76eb85f87b15328e06bd1550668f79f7b18c6/numpy-1.26.4-cp312-cp312-win32.whl", hash = "sha256:50193e430acfc1346175fcbdaa28ffec49947a06918b7b92130744e81e640110", size = 5677803, upload-time = "2024-02-05T23:58:08.963Z" }, + { url = "https://files.pythonhosted.org/packages/16/2e/86f24451c2d530c88daf997cb8d6ac622c1d40d19f5a031ed68a4b73a374/numpy-1.26.4-cp312-cp312-win_amd64.whl", hash = "sha256:08beddf13648eb95f8d867350f6a018a4be2e5ad54c8d8caed89ebca558b2818", size = 15517754, upload-time = "2024-02-05T23:58:36.364Z" }, +] + +[[package]] +name = "nvidia-cublas-cu12" +version = "12.9.2.10" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cuda-nvrtc-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/f7/a2/c96163a0fff1839c0c9548bbdeae7b853b867009e33b9b9264adc238b1cf/nvidia_cublas_cu12-12.9.2.10-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:5572131a59c3eebeeb1c4c8144f772d49372c20124916e072a0e3fc30df421d5", size = 575012079, upload-time = "2026-04-08T18:51:47.303Z" }, + { url = "https://files.pythonhosted.org/packages/cb/c0/0a517bfe63ccd3b92eb254d264e28fca3c7cab75d07daea315250fb1bf73/nvidia_cublas_cu12-12.9.2.10-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:e4f53a8ca8c5d6e8c492d0d0a3d565ecb59a751b19cfdaa4f6da0ab2104c1702", size = 581240110, upload-time = "2026-04-08T18:52:31.532Z" }, + { url = "https://files.pythonhosted.org/packages/20/e2/fc9a0e985249d873150276d5afb02e39a66817fedbf1a385724393e505ed/nvidia_cublas_cu12-12.9.2.10-py3-none-win_amd64.whl", hash = "sha256:623f43027d40d44ceadf0043f002bd25cf353e8f13ce90b9a87057019f560661", size = 553162896, upload-time = "2026-04-08T18:53:10.035Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.9.79" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b4/78/351b5c8cdbd9a6b4fb0d6ee73fb176dcdc1b6b6ad47c2ffff5ae8ca4a1f7/nvidia_cuda_cupti_cu12-12.9.79-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:791853b030602c6a11d08b5578edfb957cadea06e9d3b26adbf8d036135a4afe", size = 10077166, upload-time = "2025-06-05T20:01:01.385Z" }, + { url = "https://files.pythonhosted.org/packages/c1/2e/b84e32197e33f39907b455b83395a017e697c07a449a2b15fd07fc1c9981/nvidia_cuda_cupti_cu12-12.9.79-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:096bcf334f13e1984ba36685ad4c1d6347db214de03dbb6eebb237b41d9d934f", size = 10814997, upload-time = "2025-06-05T20:01:10.168Z" }, + { url = "https://files.pythonhosted.org/packages/3b/b4/298983ab1a83de500f77d0add86d16d63b19d1a82c59f8eaf04f90445703/nvidia_cuda_cupti_cu12-12.9.79-py3-none-win_amd64.whl", hash = "sha256:1848a9380067560d5bee10ed240eecc22991713e672c0515f9c3d9396adf93c8", size = 7730496, upload-time = "2025-06-05T20:11:26.444Z" }, +] + +[[package]] +name = "nvidia-cuda-nvcc-cu12" +version = "12.9.86" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/25/48/b54a06168a2190572a312bfe4ce443687773eb61367ced31e064953dd2f7/nvidia_cuda_nvcc_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:5d6a0d32fdc7ea39917c20065614ae93add6f577d840233237ff08e9a38f58f0", size = 40546229, upload-time = "2025-06-05T20:01:53.357Z" }, + { url = "https://files.pythonhosted.org/packages/d6/5c/8cc072436787104bbbcbde1f76ab4a0d89e68f7cebc758dd2ad7913a43d0/nvidia_cuda_nvcc_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:44e1eca4d08926193a558d2434b1bf83d57b4d5743e0c431c0c83d51da1df62b", size = 39411138, upload-time = "2025-06-05T20:01:43.182Z" }, + { url = "https://files.pythonhosted.org/packages/d2/9e/c71c53655a65d7531c89421c282359e2f626838762f1ce6180ea0bbebd29/nvidia_cuda_nvcc_cu12-12.9.86-py3-none-win_amd64.whl", hash = "sha256:8ed7f0b17dea662755395be029376db3b94fed5cbb17c2d35cc866c5b1b84099", size = 34669845, upload-time = "2025-06-05T20:11:56.308Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.9.86" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b8/85/e4af82cc9202023862090bfca4ea827d533329e925c758f0cde964cb54b7/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:210cf05005a447e29214e9ce50851e83fc5f4358df8b453155d5e1918094dcb4", size = 89568129, upload-time = "2025-06-05T20:02:41.973Z" }, + { url = "https://files.pythonhosted.org/packages/64/eb/c2295044b8f3b3b08860e2f6a912b702fc92568a167259df5dddb78f325e/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:096d4de6bda726415dfaf3198d4f5c522b8e70139c97feef5cd2ca6d4cd9cead", size = 44528905, upload-time = "2025-06-05T20:02:29.754Z" }, + { url = "https://files.pythonhosted.org/packages/52/de/823919be3b9d0ccbf1f784035423c5f18f4267fb0123558d58b813c6ec86/nvidia_cuda_nvrtc_cu12-12.9.86-py3-none-win_amd64.whl", hash = "sha256:72972ebdcf504d69462d3bcd67e7b81edd25d0fb85a2c46d3ea3517666636349", size = 76408187, upload-time = "2025-06-05T20:12:27.819Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.9.79" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/e0/0279bd94539fda525e0c8538db29b72a5a8495b0c12173113471d28bce78/nvidia_cuda_runtime_cu12-12.9.79-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:83469a846206f2a733db0c42e223589ab62fd2fabac4432d2f8802de4bded0a4", size = 3515012, upload-time = "2025-06-05T20:00:35.519Z" }, + { url = "https://files.pythonhosted.org/packages/bc/46/a92db19b8309581092a3add7e6fceb4c301a3fd233969856a8cbf042cd3c/nvidia_cuda_runtime_cu12-12.9.79-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:25bba2dfb01d48a9b59ca474a1ac43c6ebf7011f1b0b8cc44f54eb6ac48a96c3", size = 3493179, upload-time = "2025-06-05T20:00:53.735Z" }, + { url = "https://files.pythonhosted.org/packages/59/df/e7c3a360be4f7b93cee39271b792669baeb3846c58a4df6dfcf187a7ffab/nvidia_cuda_runtime_cu12-12.9.79-py3-none-win_amd64.whl", hash = "sha256:8e018af8fa02363876860388bd10ccb89eb9ab8fb0aa749aaf58430a9f7c4891", size = 3591604, upload-time = "2025-06-05T20:11:17.036Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.27.0.42" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/e4/4b/26e6309b205085e9aa8074887e81fa3603da24c51fd75e0dcaaaff0b9e0f/nvidia_cudnn_cu12-9.27.0.42-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:5bc22243563107ea06aa60f9116ad800063d5f4617cc75bc6670d5db78d9f336", size = 772694764, upload-time = "2026-09-29T23:39:53.164Z" }, + { url = "https://files.pythonhosted.org/packages/65/e4/c5a205d48ff00ed8b27882bb45d338d8138e976c76385f328a326bbfaeda/nvidia_cudnn_cu12-9.27.0.42-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0a4aa3a7d2264256506c6857fc41fc0c499982f78d70195b1cfcc9055c1957cd", size = 766178381, upload-time = "2026-09-29T23:42:52.357Z" }, + { url = "https://files.pythonhosted.org/packages/aa/38/f856579877f7c1c5066e61182e7de7bc27bf35a78c8d1b0fa592e6985bc4/nvidia_cudnn_cu12-9.27.0.42-py3-none-win_amd64.whl", hash = "sha256:06e9b0026f3bad97d2b58666330fabec04fe1672f776661ecb0ce0029c27f142", size = 743068852, upload-time = "2026-09-29T23:45:50.385Z" }, +] + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.4.1.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/9b/2b/76445b0af890da61b501fde30650a1a4bd910607261b209cccb5235d3daa/nvidia_cufft_cu12-11.4.1.4-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1a28c9b12260a1aa7a8fd12f5ebd82d027963d635ba82ff39a1acfa7c4c0fbcf", size = 200822453, upload-time = "2025-06-05T20:05:27.889Z" }, + { url = "https://files.pythonhosted.org/packages/95/f4/61e6996dd20481ee834f57a8e9dca28b1869366a135e0d42e2aa8493bdd4/nvidia_cufft_cu12-11.4.1.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c67884f2a7d276b4b80eb56a79322a95df592ae5e765cf1243693365ccab4e28", size = 200877592, upload-time = "2025-06-05T20:05:45.862Z" }, + { url = "https://files.pythonhosted.org/packages/20/ee/29955203338515b940bd4f60ffdbc073428f25ef9bfbce44c9a066aedc5c/nvidia_cufft_cu12-11.4.1.4-py3-none-win_amd64.whl", hash = "sha256:8e5bfaac795e93f80611f807d42844e8e27e340e0cde270dcb6c65386d795b80", size = 200067309, upload-time = "2025-06-05T20:13:59.762Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.7.5.82" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12" }, + { name = "nvidia-cusparse-cu12" }, + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/03/99/686ff9bf3a82a531c62b1a5c614476e8dfa24a9d89067aeedf3592ee4538/nvidia_cusolver_cu12-11.7.5.82-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:62efa83e4ace59a4c734d052bb72158e888aa7b770e1a5f601682f16fe5b4fd2", size = 337869834, upload-time = "2025-06-05T20:06:53.125Z" }, + { url = "https://files.pythonhosted.org/packages/33/40/79b0c64d44d6c166c0964ec1d803d067f4a145cca23e23925fd351d0e642/nvidia_cusolver_cu12-11.7.5.82-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:15da72d1340d29b5b3cf3fd100e3cd53421dde36002eda6ed93811af63c40d88", size = 338117415, upload-time = "2025-06-05T20:07:16.809Z" }, + { url = "https://files.pythonhosted.org/packages/32/5d/feb7f86b809f89b14193beffebe24cf2e4bf7af08372ab8cdd34d19a65a0/nvidia_cusolver_cu12-11.7.5.82-py3-none-win_amd64.whl", hash = "sha256:77666337237716783c6269a658dea310195cddbd80a5b2919b1ba8735cec8efd", size = 326215953, upload-time = "2025-06-05T20:14:41.76Z" }, +] + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.5.10.65" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/6f/8710fbd17cdd1d0fc3fea7d36d5b65ce1933611c31e1861da330206b253a/nvidia_cusparse_cu12-12.5.10.65-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:221c73e7482dd93eda44e65ce567c031c07e2f93f6fa0ecd3ba876a195023e83", size = 366359408, upload-time = "2025-06-05T20:07:42.501Z" }, + { url = "https://files.pythonhosted.org/packages/12/46/b0fd4b04f86577921feb97d8e2cf028afe04f614d17fb5013de9282c9216/nvidia_cusparse_cu12-12.5.10.65-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:73060ce019ac064a057267c585bf1fd5a353734151f87472ff02b2c5c9984e78", size = 366465088, upload-time = "2025-06-05T20:08:20.413Z" }, + { url = "https://files.pythonhosted.org/packages/73/ef/063500c25670fbd1cbb0cd3eb7c8a061585b53adb4dd8bf3492bb49b0df3/nvidia_cusparse_cu12-12.5.10.65-py3-none-win_amd64.whl", hash = "sha256:9e487468a22a1eaf1fbd1d2035936a905feb79c4ce5c2f67626764ee4f90227c", size = 362504719, upload-time = "2025-06-05T20:15:17.947Z" }, +] + +[[package]] +name = "nvidia-nccl-cu12" +version = "2.32.3" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/cb/702524922e4f64e42e8c299a26c7d794bce50dee2f5f54758f4d1e162305/nvidia_nccl_cu12-2.32.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:061af42ae1044816820e16b872075e243d1ac0656f51860fbc8e106c257aa83a", size = 351519937, upload-time = "2026-09-22T08:31:48.829Z" }, + { url = "https://files.pythonhosted.org/packages/f3/f8/9dce3698eed6fae28078eb38f748feb1665767693d8017c099de9fbf0b08/nvidia_nccl_cu12-2.32.3-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:bb94b2348e4861d5e8aa00b36c3d12d9bdf9c6369b82cbb723de0b582280a21c", size = 351526311, upload-time = "2026-09-22T08:32:27.013Z" }, +] + +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.9.86" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/0c/c75bbfb967457a0b7670b8ad267bfc4fffdf341c074e0a80db06c24ccfd4/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:e3f1171dbdc83c5932a45f0f4c99180a70de9bd2718c1ab77d14104f6d7147f9", size = 39748338, upload-time = "2025-06-05T20:10:25.613Z" }, + { url = "https://files.pythonhosted.org/packages/97/bc/2dcba8e70cf3115b400fef54f213bcd6715a3195eba000f8330f11e40c45/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:994a05ef08ef4b0b299829cde613a424382aff7efb08a7172c1fa616cc3af2ca", size = 39514880, upload-time = "2025-06-05T20:10:04.89Z" }, + { url = "https://files.pythonhosted.org/packages/dd/7e/2eecb277d8a98184d881fb98a738363fd4f14577a4d2d7f8264266e82623/nvidia_nvjitlink_cu12-12.9.86-py3-none-win_amd64.whl", hash = "sha256:cc6fcec260ca843c10e34c936921a1c426b351753587fdd638e8cff7b16bb9db", size = 35584936, upload-time = "2025-06-05T20:16:08.525Z" }, +] + +[[package]] +name = "opt-einsum" +version = "3.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8c/b9/2ac072041e899a52f20cf9510850ff58295003aa75525e58343591b0cbfb/opt_einsum-3.4.0.tar.gz", hash = "sha256:96ca72f1b886d148241348783498194c577fa30a8faac108586b14f1ba4473ac", size = 63004, upload-time = "2024-09-26T14:33:24.483Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl", hash = "sha256:69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd", size = 71932, upload-time = "2024-09-26T14:33:23.039Z" }, +] + +[[package]] +name = "optax" +version = "0.2.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "chex" }, + { name = "jax" }, + { name = "jaxlib" }, + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c0/75/1e011953c48be502d4d84fa8458e91be7c6f983002511669bddd7b1a065f/optax-0.2.5.tar.gz", hash = "sha256:b2e38c7aea376186deae758ba7a258e6ef760c6f6131e9e11bc561c65386d594", size = 258548, upload-time = "2025-06-10T17:00:47.544Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b9/33/f86091c706db1a5459f501830241afff2ecab3532725c188ea57be6e54de/optax-0.2.5-py3-none-any.whl", hash = "sha256:966deae936207f268ac8f564d8ed228d645ac1aaddefbbf194096d2299b24ba8", size = 354324, upload-time = "2025-06-10T17:00:46.062Z" }, +] + +[[package]] +name = "optree" +version = "0.20.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/76/e2/89ef1e5ef78ffd22c1d1e66d884c33850e051a11780209c35503a02855a3/optree-0.20.0.tar.gz", hash = "sha256:c7403eb0f2b2a060a97803a8f52904212df85ed6cff4b0eeed9cc6e38c2cf88d", size = 244699, upload-time = "2026-08-20T11:38:46.956Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/27/6d/a7863d089b6d305ce413e4c21cabe4db83306959a7442e77430a33d9c285/optree-0.20.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:5e81ba65acc15054b6a4f97522a7a970daeb0c6e9c066563552b9542b565bde3", size = 450795, upload-time = "2026-08-20T11:36:31.222Z" }, + { url = "https://files.pythonhosted.org/packages/cc/b3/be34d1ca24f842a626dc5f7e4353328f7ffc341c54b46e61457ec54b4a1f/optree-0.20.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:05596b84c1d5b43d9f45ca0bfee528e3e96f0d547522b074cf20bfe8ffea8938", size = 417675, upload-time = "2026-08-20T11:36:32.961Z" }, + { url = "https://files.pythonhosted.org/packages/2f/c0/be1c52ebb094637b235d67b07a7c852270ef24a34ec3a653d97fcebc91b4/optree-0.20.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7ce0418dac7763358e96bacc7d586cfc5b820af2c7c17d1e48331edc806e1f12", size = 439855, upload-time = "2026-08-20T11:36:34.3Z" }, + { url = "https://files.pythonhosted.org/packages/2f/5a/7213d97897ebc6a3f51111053faff85a4f21acc941edd774f3ad1fea567d/optree-0.20.0-cp312-cp312-manylinux_2_26_i686.manylinux_2_28_i686.whl", hash = "sha256:ef1af8e371d21cc17da28e2412471880b4f7ef53c9ffa2d7452b8d021f7199fc", size = 509278, upload-time = "2026-08-20T11:36:35.477Z" }, + { url = "https://files.pythonhosted.org/packages/7f/aa/6036c15a1a578da66197d23ee3068b0058531ef3165403d99a41bb7d5ec3/optree-0.20.0-cp312-cp312-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:8ba63097249199e93fe466e6e9f68892c34e1643437a8b1ec333c49b810cae1e", size = 510560, upload-time = "2026-08-20T11:36:36.776Z" }, + { url = "https://files.pythonhosted.org/packages/14/ff/b4283cd1f78af6aae60bc50030d41631ac5fd67e3c9304368f2740e46265/optree-0.20.0-cp312-cp312-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:e79f0e3d19b9f26e3733e8c5a2669802bc9f9d79921b8f5668ea713bafc0cd1b", size = 505476, upload-time = "2026-08-20T11:36:38.121Z" }, + { url = "https://files.pythonhosted.org/packages/fe/98/6a63875d66bef096d0c040f0cba9d0b14436df88b130d7f779c3dfa59759/optree-0.20.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5c4ed341ea1ee7eb04eff01527b88a868cb19aac672a667e0aefadf7d14c0260", size = 488981, upload-time = "2026-08-20T11:36:39.535Z" }, + { url = "https://files.pythonhosted.org/packages/08/df/b61d57cf0b40b0a1908e3170cd74b36b4425938a8e9ff6b97005d982b312/optree-0.20.0-cp312-cp312-manylinux_2_39_riscv64.whl", hash = "sha256:40659c9c055e4b0be8f9ac9951fe1280199fd30d94940ac619da9ee010d1c67e", size = 450539, upload-time = "2026-08-20T11:36:40.762Z" }, + { url = "https://files.pythonhosted.org/packages/cc/de/7d78a30503f5f44c9a47a9b76f13f15d41bd38f1c9f7a194285aa59013cc/optree-0.20.0-cp312-cp312-win32.whl", hash = "sha256:ac4077d1a655edfe5fbb92d1be79afc5fdab9f0f244586c391231db276f3cc9d", size = 528162, upload-time = "2026-08-20T11:36:41.864Z" }, + { url = "https://files.pythonhosted.org/packages/5c/79/c7af080d3301a4912f39d99faad277765d2ebf7f2faea3d2519837247623/optree-0.20.0-cp312-cp312-win_amd64.whl", hash = "sha256:ed117076eab16f8ce4510efffcc81b08f0f50cbbc14f3b8b550aa46e1ebc97b1", size = 567012, upload-time = "2026-08-20T11:36:42.99Z" }, + { url = "https://files.pythonhosted.org/packages/2a/74/4dd545f53e77ec92fed816dd70c9735583b4cb035bcf6aaf0f5914f8eb05/optree-0.20.0-cp312-cp312-win_arm64.whl", hash = "sha256:c60b206a42a3225fa8b9a2a71d1b0b7a0659427c6deace3b742b925f16371cb0", size = 739199, upload-time = "2026-08-20T11:36:44.143Z" }, +] + +[[package]] +name = "orbax-checkpoint" +version = "0.6.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "etils", extra = ["epath", "epy"] }, + { name = "humanize" }, + { name = "jax" }, + { name = "jaxlib" }, + { name = "msgpack" }, + { name = "nest-asyncio" }, + { name = "numpy" }, + { name = "protobuf" }, + { name = "pyyaml" }, + { name = "tensorstore" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/32/9f/6ca67c33b6a53d2484ebeb3c132a06331fbf085ce1619e3f919d36bb6a9c/orbax_checkpoint-0.6.4.tar.gz", hash = "sha256:366b4d528a7322e1b3d9ddcaed45c8515add0d2fc69c8975c30d98638543240f", size = 194820, upload-time = "2024-09-17T17:37:10.151Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/79/0a84020557a5c5bd987f74892e58e71cf74886dc9b264ad7630a1f288543/orbax_checkpoint-0.6.4-py3-none-any.whl", hash = "sha256:b4f2608ee4d1da67f7619fc35ff9c928ecdf4ccf7546eeb43ecf38c2608b6dea", size = 270502, upload-time = "2024-09-17T17:37:07.477Z" }, +] + +[[package]] +name = "packaging" +version = "26.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/fa/3944b40b07da9ce895c0e6303a5ab7d53da063554f534556b134a54d6093/packaging-26.3.tar.gz", hash = "sha256:94edc256424af38762eb31306eed28beb9f0efc50a8837492c9d6fd6004aed79", size = 313412, upload-time = "2026-08-04T18:15:28.737Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/34/ba1c580383c9eada3711951fef0795c80b829a078d72188184bcab9dd527/packaging-26.3-py3-none-any.whl", hash = "sha256:d7193f7c8e4e93f444fde0262bf90af30e16fa0ad0ad44cb553c87339b23cd1c", size = 129956, upload-time = "2026-08-04T18:15:27.159Z" }, +] + +[[package]] +name = "pillow" +version = "13.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ce/a4/9cc30690a394a8305e8be4c9036aa16d0dfddadd09079860b582f562496d/pillow-13.0.0.tar.gz", hash = "sha256:cdcd1ba258178f7356d74170c9a81ba4eab6a5d55748bc402a5a2a6fe6be3510", size = 47058631, upload-time = "2026-10-16T22:18:09.718Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7a/b8/58d79332229482eb0d27c82973850808c19d37e5bdc14443d3417bd373f5/pillow-13.0.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:67c8d44ad89fa7eb88de5d6ef6199b1e03dd24974f85a67639fc8badb801fe79", size = 5568643, upload-time = "2026-10-16T22:15:44.474Z" }, + { url = "https://files.pythonhosted.org/packages/39/4b/f2ed53d83aa44ac19599ac0a54897f867fafb95f406980d1e24083f8cbe3/pillow-13.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:fc481f742ccda20e9ed760b8f978dff50f2a4e9bc253ac167fb0a05b0783ef52", size = 4852329, upload-time = "2026-10-16T22:15:46.284Z" }, + { url = "https://files.pythonhosted.org/packages/3b/bb/9e30e8ed18c20800ee327a12107d5249ad8dc444867b3adfcc21654e3c71/pillow-13.0.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8bf295fe459343bc6093c8e5e0df60bfa189e0fcb775a22e2c2488fdeccc6c0d", size = 6519309, upload-time = "2026-10-16T22:15:48.14Z" }, + { url = "https://files.pythonhosted.org/packages/e6/99/fdf65cc539e266ab93546c261bdf98411fe165c6e4e13f9dbf597097192b/pillow-13.0.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:42d63be3af4f9b951b199b0a9d23365d37ab14687e6d77b74940432c8515c4ca", size = 7144542, upload-time = "2026-10-16T22:15:50.093Z" }, + { url = "https://files.pythonhosted.org/packages/26/9b/f38b7c4d8f599537e2d1a6d546cec9beef7536ff8f5c49a3f9c78179d8ac/pillow-13.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:4bd024d567eca931b668978d672dba914d4b8e2628e2c3073a2ad178de8a2620", size = 6618042, upload-time = "2026-10-16T22:15:52.054Z" }, + { url = "https://files.pythonhosted.org/packages/a8/7c/45ee28531f89ee4582967d8a17f7a26d71ca64cd902b93c269590cd60a00/pillow-13.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:23edf4ec4b6f4f240e8df6e4985e9ee2ced86474632433e55fb1e2548744bf40", size = 7228137, upload-time = "2026-10-16T22:15:54.009Z" }, + { url = "https://files.pythonhosted.org/packages/01/bc/9c2a969cb1b70201d04fa17fd5aedf02a732870b8da86bceaaa30b0ca385/pillow-13.0.0-cp312-cp312-win32.whl", hash = "sha256:02464da86f561f410e13b061ec37c3068a9cb557a095494157742d07b3c24a27", size = 6515423, upload-time = "2026-10-16T22:15:56.25Z" }, + { url = "https://files.pythonhosted.org/packages/ff/cd/73a79ecdc822076002eb24b5e1e4a530b0bb5310fcb6d239abbec877110b/pillow-13.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:ff3fe9eb621fbb9db4c4be33ab946f6bac5f98ab50ca8d3260503653656c961b", size = 7272166, upload-time = "2026-10-16T22:15:58.597Z" }, + { url = "https://files.pythonhosted.org/packages/4a/a3/be5a1236efd085b3984e4305da8f179ae9dab80c991d2ad78f2bd5c4c843/pillow-13.0.0-cp312-cp312-win_arm64.whl", hash = "sha256:c3bc8ee0271387e0b1e3ad0ff893645990a656e79bf9c42f4f734a3642b11d51", size = 2698943, upload-time = "2026-10-16T22:16:00.637Z" }, + { url = "https://files.pythonhosted.org/packages/60/da/f6d644e631c5b370c5d0b899f3ccadd2780aab070b29a110527cf4e1fd2c/pillow-13.0.0-pp312-pypy312_pp80-macosx_10_15_x86_64.whl", hash = "sha256:8619685cad8efc4c38c66d3b36aad854523297c43d748b3b6f0cc03c90ba3463", size = 5523188, upload-time = "2026-10-16T22:17:59.334Z" }, + { url = "https://files.pythonhosted.org/packages/e4/b0/2594f202d743a35d4fe80682bfd96f277e5ca7aecb27cf64ef5bfc1d000b/pillow-13.0.0-pp312-pypy312_pp80-macosx_11_0_arm64.whl", hash = "sha256:7e91e97a929232cdb828889ee6ad0e48d7f40f4ec38818f4e4a8968b2b1af22e", size = 4806808, upload-time = "2026-10-16T22:18:01.188Z" }, + { url = "https://files.pythonhosted.org/packages/8e/ed/d504165ef79af6fafe25c76726cf30f544cc236002f7ed054f1d51adc4ff/pillow-13.0.0-pp312-pypy312_pp80-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:93ff9c9c3663ce23aaa5d0b6b6abac55359e2a0c06c7e30f3771829af816be63", size = 5464358, upload-time = "2026-10-16T22:18:02.779Z" }, + { url = "https://files.pythonhosted.org/packages/7b/fe/a10aba5503ca13e9e009b8bd0a1a4ce5584bdce246245fd6f3a83701369c/pillow-13.0.0-pp312-pypy312_pp80-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e65dd6e75de67b2362a4fa3e16aa627595f65054dc4ec73d64680c84d592fa61", size = 6043514, upload-time = "2026-10-16T22:18:04.538Z" }, + { url = "https://files.pythonhosted.org/packages/c0/37/66b12dc476041fe523bca29c9d58fbf1ea92a4eaf73acfb9e993948f1191/pillow-13.0.0-pp312-pypy312_pp80-win_amd64.whl", hash = "sha256:567f20f1c75897d7663830da6ea8af719e03d712b1d9ee7f813319f9ece42e5a", size = 7366030, upload-time = "2026-10-16T22:18:06.638Z" }, +] + +[[package]] +name = "pluggy" +version = "1.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/bf/db/7fc19e6f2dc92a966727031389fc2e08b558f0f25eb7403c1119ad4713cd/pluggy-1.7.0.tar.gz", hash = "sha256:d1eaa46ebb595891b860ab086b4d09c8588af65ebd4361b8e8f4bb8920b90ba8", size = 123304, upload-time = "2026-10-15T09:50:58.343Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/40/9e/2b38731e0fc536806f16490e1a12d7f0dc2a1235aa8cc07bcc75416a7daa/pluggy-1.7.0-py3-none-any.whl", hash = "sha256:7dd7b0d8832ba3cb632c306926ded123429211b83641b35dc5c41ad2d34f9bec", size = 27082, upload-time = "2026-10-15T09:50:56.808Z" }, +] + +[[package]] +name = "protobuf" +version = "4.25.9" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d9/8e/d08c41a8c004e1d437ef467e7c4f9c3295cd784eba48ed5d1d01f94b1dad/protobuf-4.25.9.tar.gz", hash = "sha256:b0dc7e7c68de8b1ce831dacb12fb407e838edbb8b6cc0dc3a2a6b4cbf6de9cff", size = 381040, upload-time = "2026-03-25T23:09:36.423Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a8/e9/59435bd04bdd46cb38c42a336b22f9843e8e586ff83c35a5423f8b14704e/protobuf-4.25.9-cp310-abi3-win32.whl", hash = "sha256:bde396f568b0b46fc8fbfe9f02facf25b6755b2578a3b8ac61e74b9d69499e03", size = 392879, upload-time = "2026-03-25T23:09:21.32Z" }, + { url = "https://files.pythonhosted.org/packages/f3/16/42a5c7f1001783d2b5bfcecde10127f09010f78982c86ae409122ce3ece6/protobuf-4.25.9-cp310-abi3-win_amd64.whl", hash = "sha256:3683c05154252206f7cb2d371626514b3708199d9bcf683b503dabf3a2e38e06", size = 413900, upload-time = "2026-03-25T23:09:23.589Z" }, + { url = "https://files.pythonhosted.org/packages/56/5b/0074a0a9eb01f3d1c4648ca5e81b22090c811b210b61df9018ac6d6c5cda/protobuf-4.25.9-cp37-abi3-macosx_10_9_universal2.whl", hash = "sha256:9560813560e6ee72c11ca8873878bdb7ee003c96a57ebb013245fe84e2540904", size = 394826, upload-time = "2026-03-25T23:09:25.194Z" }, + { url = "https://files.pythonhosted.org/packages/54/aa/b2dba856f64c36b2a06c67be1472de98cca07a2322d0f0cbf03279a40e5b/protobuf-4.25.9-cp37-abi3-manylinux2014_aarch64.whl", hash = "sha256:999146ef02e7fa6a692477badd1528bcd7268df211852a3df2d834ba2b480791", size = 294191, upload-time = "2026-03-25T23:09:26.613Z" }, + { url = "https://files.pythonhosted.org/packages/a8/5c/53f18822017b8bda6bd8bb4e02048e911fdc79a3dafdc83ab994fe922a84/protobuf-4.25.9-cp37-abi3-manylinux2014_x86_64.whl", hash = "sha256:438c636de8fb706a0de94a12a268ef1ae8f5ba5ae655a7671fcda5968ba3c9be", size = 295178, upload-time = "2026-03-25T23:09:27.839Z" }, + { url = "https://files.pythonhosted.org/packages/16/28/d5065b212685875d3924bcdb3201cbf467cb4d58a18aa19a8dfd99ea80a9/protobuf-4.25.9-py3-none-any.whl", hash = "sha256:d49b615e7c935194ac161f0965699ac84df6112c378e05ec53da65d2e4cbb6d4", size = 156822, upload-time = "2026-03-25T23:09:34.957Z" }, +] + +[[package]] +name = "pygments" +version = "2.21.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/49/2e/ced460408999b33da6b31b0021b0f37d329e202d4169aeb164493778f25b/pygments-2.21.0.tar.gz", hash = "sha256:610ca751c9bc2492b38eb9a38a7fbc93edbbb2d7182edaf34e66ae493dee5c8c", size = 5005329, upload-time = "2026-08-17T08:02:48.824Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/46/17f022dd3e953bf20a04a028a21ec746d942f8d2af30fa0f124fa0e6a684/pygments-2.21.0-py3-none-any.whl", hash = "sha256:2363c69b61c4a97c838da3b130dcd6468f4848992b21a82f2a63ec34377137d9", size = 1250147, upload-time = "2026-08-17T08:02:44.912Z" }, +] + +[[package]] +name = "pyopengl" +version = "3.1.10" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6f/16/912b7225d56284859cd9a672827f18be43f8012f8b7b932bc4bd959a298e/pyopengl-3.1.10.tar.gz", hash = "sha256:c4a02d6866b54eb119c8e9b3fb04fa835a95ab802dd96607ab4cdb0012df8335", size = 1915580, upload-time = "2025-08-18T02:33:01.76Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/de/e4/1ba6f44e491c4eece978685230dde56b14d51a0365bc1b774ddaa94d14cd/pyopengl-3.1.10-py3-none-any.whl", hash = "sha256:794a943daced39300879e4e47bd94525280685f42dbb5a998d336cfff151d74f", size = 3194996, upload-time = "2025-08-18T02:32:59.902Z" }, +] + +[[package]] +name = "pytest" +version = "9.2.0" +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/79/42/bb4641652e58dedeac2ec51abc1f8cae4fee4ec63899a0999016229a0f2a/pytest-9.2.0.tar.gz", hash = "sha256:97c4f3b1f3c7414aa36465e09b8220f1d03627993bd51b1282e193753c80557b", size = 1797555, upload-time = "2026-10-27T21:30:20.728Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/51/39/1f5d8c8877a0fa327db959c19341d4be61660de1ab6b7391f9eb805bcd67/pytest-9.2.0-py3-none-any.whl", hash = "sha256:07acdcf88ec174f907dc01af70bf52d49b014bbab03faeeade31c96b18c2e588", size = 411776, upload-time = "2026-10-27T21:30:19.131Z" }, +] + +[[package]] +name = "pyyaml" +version = "6.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/05/8e/961c0007c59b8dd7729d542c61a4d537767a59645b82a0b521206e1e25c2/pyyaml-6.0.3.tar.gz", hash = "sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f", size = 130960, upload-time = "2025-09-25T21:33:16.546Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/33/422b98d2195232ca1826284a76852ad5a86fe23e31b009c9886b2d0fb8b2/pyyaml-6.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7f047e29dcae44602496db43be01ad42fc6f1cc0d8cd6c83d342306c32270196", size = 182063, upload-time = "2025-09-25T21:32:11.445Z" }, + { url = "https://files.pythonhosted.org/packages/89/a0/6cf41a19a1f2f3feab0e9c0b74134aa2ce6849093d5517a0c550fe37a648/pyyaml-6.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:fc09d0aa354569bc501d4e787133afc08552722d3ab34836a80547331bb5d4a0", size = 173973, upload-time = "2025-09-25T21:32:12.492Z" }, + { url = "https://files.pythonhosted.org/packages/ed/23/7a778b6bd0b9a8039df8b1b1d80e2e2ad78aa04171592c8a5c43a56a6af4/pyyaml-6.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9149cad251584d5fb4981be1ecde53a1ca46c891a79788c0df828d2f166bda28", size = 775116, upload-time = "2025-09-25T21:32:13.652Z" }, + { url = "https://files.pythonhosted.org/packages/65/30/d7353c338e12baef4ecc1b09e877c1970bd3382789c159b4f89d6a70dc09/pyyaml-6.0.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5fdec68f91a0c6739b380c83b951e2c72ac0197ace422360e6d5a959d8d97b2c", size = 844011, upload-time = "2025-09-25T21:32:15.21Z" }, + { url = "https://files.pythonhosted.org/packages/8b/9d/b3589d3877982d4f2329302ef98a8026e7f4443c765c46cfecc8858c6b4b/pyyaml-6.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ba1cc08a7ccde2d2ec775841541641e4548226580ab850948cbfda66a1befcdc", size = 807870, upload-time = "2025-09-25T21:32:16.431Z" }, + { url = "https://files.pythonhosted.org/packages/05/c0/b3be26a015601b822b97d9149ff8cb5ead58c66f981e04fedf4e762f4bd4/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8dc52c23056b9ddd46818a57b78404882310fb473d63f17b07d5c40421e47f8e", size = 761089, upload-time = "2025-09-25T21:32:17.56Z" }, + { url = "https://files.pythonhosted.org/packages/be/8e/98435a21d1d4b46590d5459a22d88128103f8da4c2d4cb8f14f2a96504e1/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:41715c910c881bc081f1e8872880d3c650acf13dfa8214bad49ed4cede7c34ea", size = 790181, upload-time = "2025-09-25T21:32:18.834Z" }, + { url = "https://files.pythonhosted.org/packages/74/93/7baea19427dcfbe1e5a372d81473250b379f04b1bd3c4c5ff825e2327202/pyyaml-6.0.3-cp312-cp312-win32.whl", hash = "sha256:96b533f0e99f6579b3d4d4995707cf36df9100d67e0c8303a0c55b27b5f99bc5", size = 137658, upload-time = "2025-09-25T21:32:20.209Z" }, + { url = "https://files.pythonhosted.org/packages/86/bf/899e81e4cce32febab4fb42bb97dcdf66bc135272882d1987881a4b519e9/pyyaml-6.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:5fcd34e47f6e0b794d17de1b4ff496c00986e1c83f7ab2fb8fcfe9616ff7477b", size = 154003, upload-time = "2025-09-25T21:32:21.167Z" }, + { url = "https://files.pythonhosted.org/packages/1a/08/67bd04656199bbb51dbed1439b7f27601dfb576fb864099c7ef0c3e55531/pyyaml-6.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:64386e5e707d03a7e172c0701abfb7e10f0fb753ee1d773128192742712a98fd", size = 140344, upload-time = "2025-09-25T21:32:22.617Z" }, +] + +[[package]] +name = "requests" +version = "2.34.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "charset-normalizer" }, + { name = "idna" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ac/c3/e2a2b89f2d3e2179abd6d00ebd70bff6273f37fb3e0cc209f48b39d00cbf/requests-2.34.2.tar.gz", hash = "sha256:f288924cae4e29463698d6d60bc6a4da69c89185ad1e0bcc4104f584e960b9ed", size = 142856, upload-time = "2026-05-14T19:25:27.735Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" }, +] + +[[package]] +name = "rich" +version = "15.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c0/8f/0722ca900cc807c13a6a0c696dacf35430f72e0ec571c4275d2371fca3e9/rich-15.0.0.tar.gz", hash = "sha256:edd07a4824c6b40189fb7ac9bc4c52536e9780fbbfbddf6f1e2502c31b068c36", size = 230680, upload-time = "2026-04-12T08:24:00.75Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl", hash = "sha256:33bd4ef74232fb73fe9279a257718407f169c09b78a87ad3d296f548e27de0bb", size = 310654, upload-time = "2026-04-12T08:24:02.83Z" }, +] + +[[package]] +name = "ruamel-yaml" +version = "0.19.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c7/3b/ebda527b56beb90cb7652cb1c7e4f91f48649fbcd8d2eb2fb6e77cd3329b/ruamel_yaml-0.19.1.tar.gz", hash = "sha256:53eb66cd27849eff968ebf8f0bf61f46cdac2da1d1f3576dd4ccee9b25c31993", size = 142709, upload-time = "2026-01-02T16:50:31.84Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b8/0c/51f6841f1d84f404f92463fc2b1ba0da357ca1e3db6b7fbda26956c3b82a/ruamel_yaml-0.19.1-py3-none-any.whl", hash = "sha256:27592957fedf6e0b62f281e96effd28043345e0e66001f97683aa9a40c667c93", size = 118102, upload-time = "2026-01-02T16:50:29.201Z" }, +] + +[[package]] +name = "scipy" +version = "1.17.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/35/48/b992b488d6f299dbe3f11a20b24d3dda3d46f1a635ede1c46b5b17a7b163/scipy-1.17.1-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:35c3a56d2ef83efc372eaec584314bd0ef2e2f0d2adb21c55e6ad5b344c0dcb8", size = 31610954, upload-time = "2026-02-23T00:17:49.855Z" }, + { url = "https://files.pythonhosted.org/packages/b2/02/cf107b01494c19dc100f1d0b7ac3cc08666e96ba2d64db7626066cee895e/scipy-1.17.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:fcb310ddb270a06114bb64bbe53c94926b943f5b7f0842194d585c65eb4edd76", size = 28172662, upload-time = "2026-02-23T00:18:01.64Z" }, + { url = "https://files.pythonhosted.org/packages/cf/a9/599c28631bad314d219cf9ffd40e985b24d603fc8a2f4ccc5ae8419a535b/scipy-1.17.1-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:cc90d2e9c7e5c7f1a482c9875007c095c3194b1cfedca3c2f3291cdc2bc7c086", size = 20344366, upload-time = "2026-02-23T00:18:12.015Z" }, + { url = "https://files.pythonhosted.org/packages/35/f5/906eda513271c8deb5af284e5ef0206d17a96239af79f9fa0aebfe0e36b4/scipy-1.17.1-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:c80be5ede8f3f8eded4eff73cc99a25c388ce98e555b17d31da05287015ffa5b", size = 22704017, upload-time = "2026-02-23T00:18:21.502Z" }, + { url = "https://files.pythonhosted.org/packages/da/34/16f10e3042d2f1d6b66e0428308ab52224b6a23049cb2f5c1756f713815f/scipy-1.17.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e19ebea31758fac5893a2ac360fedd00116cbb7628e650842a6691ba7ca28a21", size = 32927842, upload-time = "2026-02-23T00:18:35.367Z" }, + { url = "https://files.pythonhosted.org/packages/01/8e/1e35281b8ab6d5d72ebe9911edcdffa3f36b04ed9d51dec6dd140396e220/scipy-1.17.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:02ae3b274fde71c5e92ac4d54bc06c42d80e399fec704383dcd99b301df37458", size = 35235890, upload-time = "2026-02-23T00:18:49.188Z" }, + { url = "https://files.pythonhosted.org/packages/c5/5c/9d7f4c88bea6e0d5a4f1bc0506a53a00e9fcb198de372bfe4d3652cef482/scipy-1.17.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8a604bae87c6195d8b1045eddece0514d041604b14f2727bbc2b3020172045eb", size = 35003557, upload-time = "2026-02-23T00:18:54.74Z" }, + { url = "https://files.pythonhosted.org/packages/65/94/7698add8f276dbab7a9de9fb6b0e02fc13ee61d51c7c3f85ac28b65e1239/scipy-1.17.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f590cd684941912d10becc07325a3eeb77886fe981415660d9265c4c418d0bea", size = 37625856, upload-time = "2026-02-23T00:19:00.307Z" }, + { url = "https://files.pythonhosted.org/packages/a2/84/dc08d77fbf3d87d3ee27f6a0c6dcce1de5829a64f2eae85a0ecc1f0daa73/scipy-1.17.1-cp312-cp312-win_amd64.whl", hash = "sha256:41b71f4a3a4cab9d366cd9065b288efc4d4f3c0b37a91a8e0947fb5bd7f31d87", size = 36549682, upload-time = "2026-02-23T00:19:07.67Z" }, + { 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 = "setuptools" +version = "84.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/70/b3/c19426f783614d625699a4b247faf628ea9ab1db228674491767b417f9e1/setuptools-84.2.1.tar.gz", hash = "sha256:3ff8b3b5492a0aa66960f4419365a86b84e013070d4313e4d4a6894f5d64d6bc", size = 1230289, upload-time = "2026-10-20T19:01:37.742Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ef/55/6dc6d6bb953bdfc52519a0059592c52b6e3b8ad51cf8256dffceb55b0610/setuptools-84.2.1-py3-none-any.whl", hash = "sha256:ca140b45d5325a9478f880e351484f62fef15f336462e449ca0c03c873d589c2", size = 876263, upload-time = "2026-10-20T19:01:35.678Z" }, +] + +[[package]] +name = "six" +version = "1.17.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, +] + +[[package]] +name = "tensorboard" +version = "2.18.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "grpcio" }, + { name = "markdown" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "protobuf" }, + { name = "setuptools" }, + { name = "six" }, + { name = "tensorboard-data-server" }, + { name = "werkzeug" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/b1/de/021c1d407befb505791764ad2cbd56ceaaa53a746baed01d2e2143f05f18/tensorboard-2.18.0-py3-none-any.whl", hash = "sha256:107ca4821745f73e2aefa02c50ff70a9b694f39f790b11e6f682f7d326745eab", size = 5503036, upload-time = "2024-09-25T21:21:50.169Z" }, +] + +[[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 = "tensorflow" +version = "2.18.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "astunparse" }, + { name = "flatbuffers" }, + { name = "gast" }, + { name = "google-pasta" }, + { name = "grpcio" }, + { name = "h5py" }, + { name = "keras" }, + { name = "libclang" }, + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "opt-einsum" }, + { name = "packaging" }, + { name = "protobuf" }, + { name = "requests" }, + { name = "setuptools" }, + { name = "six" }, + { name = "tensorboard" }, + { name = "termcolor" }, + { name = "typing-extensions" }, + { name = "wrapt" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/67/8c/1cad54f8634897ad9421de8f558df9aa63d3f2747eb803ce5dbb2db1ef5b/tensorflow-2.18.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:98afa9c7f21481cdc6ccd09507a7878d533150fbb001840cc145e2132eb40942", size = 239622377, upload-time = "2025-03-12T00:13:36.89Z" }, + { url = "https://files.pythonhosted.org/packages/6c/c2/35a3542a91f4ffd4cf01e72d7f0fb59596cd5f467ff64878b0caef8e0f31/tensorflow-2.18.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1ba52b9c06ab8102b31e50acfaf56899b923171e603c8942f2bfeb181d6bb59e", size = 231996787, upload-time = "2025-03-12T00:13:47.54Z" }, + { url = "https://files.pythonhosted.org/packages/64/42/812539a8878c242eb0bacf106f5ea8936c2cc4d7f663868bd872a79772ac/tensorflow-2.18.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:442d2a774811789a8ad948e7286cb950fe3d87d3754e8cc6449d53b03dbfdaa6", size = 615623178, upload-time = "2025-03-12T00:14:03.541Z" }, + { url = "https://files.pythonhosted.org/packages/20/28/9c5e935b76eebdf46df524980d49700a9c9af56abc8c62bfd93f57709563/tensorflow-2.18.1-cp312-cp312-win_amd64.whl", hash = "sha256:210baf6d421f3e044b6e09efd04494a33b75334922fe6cf11970e2885172620a", size = 369234070, upload-time = "2025-03-12T00:14:23.423Z" }, +] + +[[package]] +name = "tensorflow-cpu" +version = "2.18.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "astunparse" }, + { name = "flatbuffers" }, + { name = "gast" }, + { name = "google-pasta" }, + { name = "grpcio" }, + { name = "h5py" }, + { name = "keras" }, + { name = "libclang" }, + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "opt-einsum" }, + { name = "packaging" }, + { name = "protobuf" }, + { name = "requests" }, + { name = "setuptools" }, + { name = "six" }, + { name = "tensorboard" }, + { name = "tensorflow-intel", marker = "sys_platform == 'win32'" }, + { name = "termcolor" }, + { name = "typing-extensions" }, + { name = "wrapt" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/ca/3f/2ed163140237aefa72c761d56af8ba3fa5cb0fe37a9f53b14ad8bcd7ef87/tensorflow_cpu-2.18.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:39bd421ad125e4163d6e2d41ab0e158b583fb5c6f9254522fb87635b0e70b891", size = 230265586, upload-time = "2024-10-25T00:17:17.656Z" }, + { url = "https://files.pythonhosted.org/packages/0e/7a/1c99bb2bb7d24238b748f9f0244a198ee15d23782bb56dbf4e7b93a29c6a/tensorflow_cpu-2.18.0-cp312-cp312-win_amd64.whl", hash = "sha256:0b093b727c2f2a8cf4ee4f2c7352c8e958a2a1d27a452961b8d5f43a0798dcd2", size = 7572, upload-time = "2024-10-25T00:17:22.893Z" }, +] + +[[package]] +name = "tensorflow-intel" +version = "2.18.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "astunparse" }, + { name = "flatbuffers" }, + { name = "gast" }, + { name = "google-pasta" }, + { name = "grpcio" }, + { name = "h5py" }, + { name = "keras" }, + { name = "libclang" }, + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "opt-einsum" }, + { name = "packaging" }, + { name = "protobuf" }, + { name = "requests" }, + { name = "setuptools" }, + { name = "six" }, + { name = "tensorboard" }, + { name = "termcolor" }, + { name = "typing-extensions" }, + { name = "wrapt" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/ae/4e/44ce609139065035c56fe570fe7f0ee8d06180c99a424bac588472052c5d/tensorflow_intel-2.18.0-cp312-cp312-win_amd64.whl", hash = "sha256:a5818043f565cf74179b67eb52fc060587ccecb9540141c39d84fbcb37ecff8c", size = 390262926, upload-time = "2024-10-25T04:46:35.62Z" }, +] + +[[package]] +name = "tensorflow-probability" +version = "0.25.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "cloudpickle" }, + { name = "decorator" }, + { name = "dm-tree" }, + { name = "gast" }, + { name = "numpy" }, + { name = "six" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/b6/e116761ceeb880b7ad4c6a7e1203538754dd7dda42538f08d6daaf8dfa4d/tensorflow_probability-0.25.0-py2.py3-none-any.whl", hash = "sha256:f3f4d6431656c0122906888afe1b67b4400e82bd7f254b45b92e6c5b84ea8e3e", size = 6979403, upload-time = "2024-11-08T16:25:57.238Z" }, +] + +[[package]] +name = "tensorstore" +version = "0.1.65" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/cb/b7/2349c1d192acc43e7e1f8752112caa206140c946b00ad0872a2939483ca0/tensorstore-0.1.65.tar.gz", hash = "sha256:65cbe5a600c32569bb0b9f597ea318cc298a13b42d5fc98168c97bb11f320eae", size = 6515607, upload-time = "2024-09-06T03:26:36.663Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c9/b3/6304f258d47c909cb5ded38749ff4f01b5996d105e15119c910f8d1526cf/tensorstore-0.1.65-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:33072ea4c19f2c102c136320cce583540fead4c4904562735a783f8acd066feb", size = 15408582, upload-time = "2024-09-06T03:26:12.616Z" }, + { url = "https://files.pythonhosted.org/packages/79/5e/4dd1b1e2698a67d1f7ee8e38c8038b4680ff391086b73172b093949a5daa/tensorstore-0.1.65-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8b64e62ce9766870f98b2cc5e98a636c728598118fc798b13d8c8d28a41f44d3", size = 13601426, upload-time = "2024-09-06T03:26:15.458Z" }, + { url = "https://files.pythonhosted.org/packages/08/da/df4bcc83be2e3c4267193f725287a5394fe5fc661e2752bf727f40b578d7/tensorstore-0.1.65-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3b4555b52a9211de1daa399a5cc4d1bdffc171cf5847bcf36f3f4dffc1957179", size = 13654691, upload-time = "2024-09-06T03:26:18.015Z" }, + { url = "https://files.pythonhosted.org/packages/cc/bd/3add7bee248e2d2177edcb501e84e081f0bf0d9b864a19037b04a7ef6c23/tensorstore-0.1.65-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:570f2c4f991a852376d0142cdcc3987ac49466a9380aebdf651d73e6700eee0f", size = 14930913, upload-time = "2024-09-06T03:26:20.445Z" }, + { url = "https://files.pythonhosted.org/packages/4c/7b/56db5850e6dea785bc0d498b43f1bbd6b9c192a0b3d4c27765d14d184b49/tensorstore-0.1.65-cp312-cp312-win_amd64.whl", hash = "sha256:5c4003abcbccbb7f4340f18ed64ebe0626a7f5cde5bf18c48fb052374049f37b", size = 11531830, upload-time = "2024-09-06T03:26:22.596Z" }, +] + +[[package]] +name = "termcolor" +version = "3.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/46/79/cf31d7a93a8fdc6aa0fbb665be84426a8c5a557d9240b6239e9e11e35fc5/termcolor-3.3.0.tar.gz", hash = "sha256:348871ca648ec6a9a983a13ab626c0acce02f515b9e1983332b17af7979521c5", size = 14434, upload-time = "2025-12-29T12:55:21.882Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/33/d1/8bb87d21e9aeb323cc03034f5eaf2c8f69841e40e4853c2627edf8111ed3/termcolor-3.3.0-py3-none-any.whl", hash = "sha256:cf642efadaf0a8ebbbf4bc7a31cec2f9b5f21a9f726f4ccbb08192c9c26f43a5", size = 7734, upload-time = "2025-12-29T12:55:20.718Z" }, +] + +[[package]] +name = "tf-keras" +version = "2.18.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "tensorflow" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a9/a4/7d0acc28cde2b29b8114552ce3258dafdc6b2186d24bf8e912de713dd74f/tf_keras-2.18.0.tar.gz", hash = "sha256:ebf744519b322afead33086a2aba872245473294affd40973694f3eb7c7ad77d", size = 1260765, upload-time = "2024-10-24T22:58:06.652Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8a/ed/e08afca471299b04a34cd548e64e89d0153eda0e6cf9b715356777e24774/tf_keras-2.18.0-py3-none-any.whl", hash = "sha256:c431d04027eef790fcd3261cf7fdf93eb74f3cb32e05078b57b7f5a54bd53262", size = 1725427, upload-time = "2024-10-24T22:58:03.918Z" }, +] + +[[package]] +name = "toolz" +version = "1.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/31/6f/ae20c212a07aa2d156c787383d8088a5e045ee39628661edb190c97e1659/toolz-1.2.0.tar.gz", hash = "sha256:9667a038e9d6ecba37995e26cb2f59ec6420b6ad8dd9677de59db9b956b08490", size = 55442, upload-time = "2026-10-07T04:16:25.639Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/db/17/4c8beb6c8c4176c6bf143bfd7e1e4dd6719b00ced90738c7ac471b71c1df/toolz-1.2.0-py3-none-any.whl", hash = "sha256:890f820b1cb8152785aaf9386d8707770110809035800985ca65cb24ce1120ef", size = 60789, upload-time = "2026-10-07T04:16:24.173Z" }, +] + +[[package]] +name = "tqdm" +version = "4.70.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0d/ea/b2a5bd54b28a324dae8211928b2d730b6547500342c7e6c6dea08bd0a485/tqdm-4.70.1.tar.gz", hash = "sha256:cefd0eca11b2a37a3aee776544d4f4ae913f02688135b5556b8788dfa474afc4", size = 171846, upload-time = "2026-09-11T07:25:16.601Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/03/921a3d3c75785aca9ebfbfcabfbc3a1be12e2ab5265deb026d55a5a3f83e/tqdm-4.70.1-py3-none-any.whl", hash = "sha256:c293e525e6fef9c20e8728fd4612df02a0aa31bb5fe91ecd93e123b1b7bffa73", size = 80199, upload-time = "2026-09-11T07:25:14.599Z" }, +] + +[[package]] +name = "treescope" +version = "0.1.10" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f0/2a/d13d3c38862632742d2fe2f7ae307c431db06538fd05ca03020d207b5dcc/treescope-0.1.10.tar.gz", hash = "sha256:20f74656f34ab2d8716715013e8163a0da79bdc2554c16d5023172c50d27ea95", size = 138870, upload-time = "2025-08-08T05:43:48.048Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/2b/36e984399089c026a6499ac8f7401d38487cf0183839a4aa78140d373771/treescope-0.1.10-py3-none-any.whl", hash = "sha256:dde52f5314f4c29d22157a6fe4d3bd103f9cae02791c9e672eefa32c9aa1da51", size = 182255, upload-time = "2025-08-08T05:43:46.673Z" }, +] + +[[package]] +name = "typing-extensions" +version = "4.16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/cc/6253133b5bb138fc3306cebfbda2c520f545d36b5be2c7255cc528bb45d6/typing_extensions-4.16.0.tar.gz", hash = "sha256:dc983d19a509c94dba722ee6abd33940f7c05a89e243c47e907eb4db6f1a43e5", size = 113555, upload-time = "2026-07-02T08:40:05.92Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/d3/b8441a820a491ddfc024b0b0cf0393375b75ea13866d9c66727e54c2fc80/typing_extensions-4.16.0-py3-none-any.whl", hash = "sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8", size = 45571, upload-time = "2026-07-02T08:40:04.659Z" }, +] + +[[package]] +name = "urllib3" +version = "2.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e3/05/b17359e1cefb4f909b5e40b1b90a496d987258916dbbf88e842c729f510e/urllib3-2.8.0.tar.gz", hash = "sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63", size = 458972, upload-time = "2026-09-15T19:29:36.253Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/92/9d/c4e665119135114480843e7ab388fa94d8480650450e6f8e26b70d323a4c/urllib3-2.8.0-py3-none-any.whl", hash = "sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3", size = 135717, upload-time = "2026-09-15T19:29:34.577Z" }, +] + +[[package]] +name = "websockets" +version = "17.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/30/4b/4692be6d5bf739ed049b3813901c42091f8c054df89d5c089343823dea0d/websockets-17.3.tar.gz", hash = "sha256:1c7dd93a1c9f89297471493674b1fc1d9467cafd45374f2172213046bdb3a329", size = 188878, upload-time = "2026-10-19T16:58:05.813Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0d/29/902ffe19c0d35603102f31490dea342f083bc77d8f8bcdc7d8b16f5810f9/websockets-17.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:41665c7ea01d3ae497a0c61fd22d3424230d18782a5fceb79f04db9ccb2042c9", size = 218250, upload-time = "2026-10-19T16:54:54.718Z" }, + { url = "https://files.pythonhosted.org/packages/f0/66/f56a2637b1aac263205b0e3cd7b2ff2e2f8b404b8118c991f4d122fd381f/websockets-17.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:8f7469408740cf06a79b01605f2ae0b2acc8839c03fcbbf05da49baee12d3302", size = 215927, upload-time = "2026-10-19T16:54:56.124Z" }, + { url = "https://files.pythonhosted.org/packages/22/7d/a27f25451c4cc62141992b51005688481210f3d3e63de3dc66f9e882ba59/websockets-17.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d68da281ac208b6236ffa8beef23fe7c7c956d8bf09b421d447232fd94441b35", size = 216196, upload-time = "2026-10-19T16:54:57.538Z" }, + { url = "https://files.pythonhosted.org/packages/9f/77/40bbc1a5c016cf9f598dea21d1250a93245c995eafaec6401d0e4eb4e647/websockets-17.3-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:078d1d41f45ee34eed5a4f3bbe7c588caecbdd6e12dfca73e8c1b0e803773793", size = 225516, upload-time = "2026-10-19T16:54:59.104Z" }, + { url = "https://files.pythonhosted.org/packages/3d/9c/91889c8c62fb7d16cd58e3074473a259303d6229bf0ce496aec87e66f3dc/websockets-17.3-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:a7f391838e8adcf1f2fe71561d9e5a9f74ea7ffa5052eb2db329aac677383136", size = 225795, upload-time = "2026-10-19T16:55:00.438Z" }, + { url = "https://files.pythonhosted.org/packages/b8/29/2e5e2158e9920f806cd45f3cbc81cf5f3eef2b9e5a74b9ef049e2a3578e3/websockets-17.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c15a165673bb3fb57facbf040edf5c0556d20ccd2e0ccc5f23ebb69fb715fdc1", size = 227034, upload-time = "2026-10-19T16:55:02.065Z" }, + { url = "https://files.pythonhosted.org/packages/b6/af/fbab5c4098b26026b877bfc139b1ea0d7c6eaf7a86b84544933b4c445f82/websockets-17.3-cp312-cp312-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:28d1935c16782a1e3650c9baa79069b016fff739e292e4cb14bbb830d6930491", size = 229863, upload-time = "2026-10-19T16:55:03.415Z" }, + { url = "https://files.pythonhosted.org/packages/ac/72/d6e460a991607794072fe3fdfcb3221a0342af61586f5bd7ea40bddcfd21/websockets-17.3-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c534c9524a2acd0b97b7022074bec3784ac273c683cf9535087e437672ea8531", size = 227669, upload-time = "2026-10-19T16:55:04.763Z" }, + { url = "https://files.pythonhosted.org/packages/ec/8f/1c7da1cfb632e141e019cf2e0bbbd4e6dc8c198c8cc29544e4acf1a89f56/websockets-17.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b459d29c8b82a420a98fa0ce8e9cf787ce0062d98cc90604ef9574fb3cc26ff8", size = 226429, upload-time = "2026-10-19T16:55:06.112Z" }, + { url = "https://files.pythonhosted.org/packages/3e/63/c48ef8bbfbd982f2cfae1cfcf9e63f5bfbb8468fd37ca700507efc0b0007/websockets-17.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:53209dffd91b0aefed9ea630fa6ad5c2944d6210ba3a8a3a8664dc1575289b14", size = 223653, upload-time = "2026-10-19T16:55:07.473Z" }, + { url = "https://files.pythonhosted.org/packages/cb/37/7bedaad48ae8e89eaa4f3942c11be0c41de77ce2118a475f3ba05e9b5a0f/websockets-17.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:9662a9fe457820c60365b3209e03cbd8e4a817195d55de65a80fa680b4bea136", size = 226465, upload-time = "2026-10-19T16:55:08.791Z" }, + { url = "https://files.pythonhosted.org/packages/33/cf/ab9a8a0de4769ca1dce392021b6a922e21b4d3687fe9e1ab4db585a088f8/websockets-17.3-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:5670a762c7ae388bf8aa04199fd436180a4d28318dca9c1eae7b38604676dbfd", size = 225071, upload-time = "2026-10-19T16:55:10.127Z" }, + { url = "https://files.pythonhosted.org/packages/89/96/64d8460cec95f873c61c6d13a09b435915150640525ae56d5871400aec09/websockets-17.3-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:ed59cb4ee258674a6ed9b8515e42079db02cb3be94f133fdd3f200183a40c49d", size = 225841, upload-time = "2026-10-19T16:55:11.519Z" }, + { url = "https://files.pythonhosted.org/packages/a6/98/2b83211ee3087db4ae66533782664c0e283cf4a08cbb14036018d0803e09/websockets-17.3-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:0f78bb2e82fe88912c9d7b25c98e57a9b08655579af89f164fd4c0657a075d62", size = 226945, upload-time = "2026-10-19T16:55:12.85Z" }, + { url = "https://files.pythonhosted.org/packages/9f/fe/6b2a37d422d7f4dfc0c92c95a14649d9bd662243ed8b0774523ed2703604/websockets-17.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:71250e010eede33c1d638e2c7a98288520d67f782fd5eceecaeca3129233a09a", size = 224502, upload-time = "2026-10-19T16:55:14.48Z" }, + { url = "https://files.pythonhosted.org/packages/0c/c5/364015c9311c89ea6085d54c70d654c5d27ba05014d44ef9fd9daa09462b/websockets-17.3-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:729989759feb81ff6493879a0d60a0bd7d2eee41001bbec0381a080318439382", size = 225483, upload-time = "2026-10-19T16:55:15.899Z" }, + { url = "https://files.pythonhosted.org/packages/c2/b3/75d0cae8398816ed34a16a78252bfc8651f2f883cc5a568e1198395f15df/websockets-17.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:90d8c4fa5ba293cf92de351fbc292cdaadb1a1d58d031a1a97ff9b8fb6250ec7", size = 225761, upload-time = "2026-10-19T16:55:17.28Z" }, + { url = "https://files.pythonhosted.org/packages/f0/af/9a7c4b2e87f47a2a6a066acdc2157e9b7cb6222a224b95a432f54ae786ce/websockets-17.3-cp312-cp312-win32.whl", hash = "sha256:0824e6ecee838cc28e6481062574c7fde49f9a2cedfacd36f0e91cb4b37fd6fe", size = 218722, upload-time = "2026-10-19T16:55:18.658Z" }, + { url = "https://files.pythonhosted.org/packages/f7/e8/e52953a8280a1e51cb9307218443f68c8f4dc4b9952db6021d5c12d03c0f/websockets-17.3-cp312-cp312-win_amd64.whl", hash = "sha256:6644eecd99c4e8cdf72e646e17b391bb70b054f44becde37dcdc48cbec7e1527", size = 219024, upload-time = "2026-10-19T16:55:19.987Z" }, + { url = "https://files.pythonhosted.org/packages/85/a4/005eb7c1f5c1a249c1231a6ed83a54ada0b6b8ffab33a608ca20e3034ea0/websockets-17.3-cp312-cp312-win_arm64.whl", hash = "sha256:4a9133dbaedc6f1af358127275676171d9a1bb3e61f47ee62b2ff9acb78af3d0", size = 218950, upload-time = "2026-10-19T16:55:22.161Z" }, + { url = "https://files.pythonhosted.org/packages/5a/1a/5582103e0362f86c4725051ded89e585ceaf750f26beace530b7e42c81c0/websockets-17.3-pp312-pypy312_pp80-macosx_10_15_x86_64.whl", hash = "sha256:0e65260a1b1d7b81651f47343d836de70b88b02358ae3658efa6382104ec37b1", size = 215872, upload-time = "2026-10-19T16:57:55.95Z" }, + { url = "https://files.pythonhosted.org/packages/ef/a7/f69d65ed3ff9ae73d428a918300c23a992e4d50683090b8dc33cc3341d5a/websockets-17.3-pp312-pypy312_pp80-macosx_11_0_arm64.whl", hash = "sha256:94c02b0702656f083c098bdbd93bce22a7aa7a2b34ba56d2472efb67dab7cf58", size = 216136, upload-time = "2026-10-19T16:57:57.251Z" }, + { url = "https://files.pythonhosted.org/packages/69/76/c64fda4ac18a3aeb32e66acec481011a766add480b6362fad80e2c3c1963/websockets-17.3-pp312-pypy312_pp80-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:ade76f8c84ade325bbed96806357fffd353fcc29a5166d6dc6d39381f6d9f76c", size = 217080, upload-time = "2026-10-19T16:57:58.593Z" }, + { url = "https://files.pythonhosted.org/packages/6b/6d/085a3f0e6152bbe79917973f3d33cd008437de6c04675e2f15253b921f21/websockets-17.3-pp312-pypy312_pp80-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:9a215749c9a88ef2842e880215193e56ae60b856773c99702b078236bc6d4adb", size = 216966, upload-time = "2026-10-19T16:57:59.924Z" }, + { url = "https://files.pythonhosted.org/packages/ed/b8/7eb904583ea2090b329f50487e6f8dc58eb558712b9f3c16f289f0be668e/websockets-17.3-pp312-pypy312_pp80-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4e7b5d23ef144b32645d97bd99f382c4394c443395f45fe1576786838d26e308", size = 217828, upload-time = "2026-10-19T16:58:01.822Z" }, + { url = "https://files.pythonhosted.org/packages/79/36/f54b4a1cd9456ca0a65a76d77d181621433793aca633dfeed214748ff4f0/websockets-17.3-pp312-pypy312_pp80-win_amd64.whl", hash = "sha256:13eace1b315345ebac5b27821f5dcce24e5cf9976ab18f79e00a59bdd67e5d19", size = 306301, upload-time = "2026-10-19T16:58:03.161Z" }, + { url = "https://files.pythonhosted.org/packages/d7/b0/3f01fcb3406662f6deb1e3e72b34dd0979e993991c0b7435a9e12bc14f02/websockets-17.3-py3-none-any.whl", hash = "sha256:dc2b495a92b3971966bf102e73589f950db3ec1ca6713265355ce7771d97b199", size = 212375, upload-time = "2026-10-19T16:58:04.483Z" }, +] + +[[package]] +name = "werkzeug" +version = "3.1.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a4/34/4dd12fc8bb7d61c91467ec3efe415ffa7d5456f799954b40c5bbaeae470e/werkzeug-3.1.9.tar.gz", hash = "sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060", size = 940188, upload-time = "2026-09-27T18:33:41.637Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a1/38/df03f564f43cec2684823f3cccae1a652ee7face1cbaa76fb223096e64d7/werkzeug-3.1.9-py3-none-any.whl", hash = "sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab", size = 228700, upload-time = "2026-09-27T18:33:39.685Z" }, +] + +[[package]] +name = "wheel" +version = "0.48.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "packaging" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d0/20/50ed6bdf27dec98b568a8ae25dc599f35baa3d9709f9e83fd1edb56b9a90/wheel-0.48.0.tar.gz", hash = "sha256:94800765601e9171bf5d58d066e640662842bcedcbab982b2c90787a2c987322", size = 66471, upload-time = "2026-08-11T22:02:27.327Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2e/29/69cfbb602cd91690c55d38ba9fe53e6a7e76a6fa647bf38f19c138d25449/wheel-0.48.0-py3-none-any.whl", hash = "sha256:3217dcc807155e45db462d7ef2431f5ddda0d7273b700d05a67b271ceb1287ab", size = 33320, upload-time = "2026-08-11T22:02:26.1Z" }, +] + +[[package]] +name = "wrapt" +version = "2.5.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9c/b0/f752d19e1e83ab8bd64aeccfde32ff086aee3ef08520f7340b69a8d0a4d2/wrapt-2.5.2.tar.gz", hash = "sha256:551cb49e88b3722faad780303833bc4fa70685fb1773270fb98b40d39d9fddbf", size = 186066, upload-time = "2026-10-29T23:06:39.868Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4b/fe/9ccae0a147d0732397f54ff232b1a0da148a4099f17c1b156b9d3e79127e/wrapt-2.5.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:bec3b6b1ae2f7ed763fb0ba1776b9b5f11283e6ff40c9cbe3a855ea652c326cf", size = 106628, upload-time = "2026-10-29T23:04:44.962Z" }, + { url = "https://files.pythonhosted.org/packages/d5/ca/a61d994df49fa422abd859e525e04b95ef2888e3818a0f451aa7d4d3527f/wrapt-2.5.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:e57d79e3338b3cecee1dd80d6ef8f945f8ea7b92f0d8d5591f840ffa0afc14dc", size = 106264, upload-time = "2026-10-29T23:04:46.307Z" }, + { url = "https://files.pythonhosted.org/packages/dd/26/a6caa9a84ef5db4a273787949cea7ab51f6001d2548768af4c8e3c49debd/wrapt-2.5.2-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:092472b7b1494deef40dac5e9f6295013a9d51b112b5fbc4350f64dc2574e46a", size = 250614, upload-time = "2026-10-29T23:04:47.735Z" }, + { url = "https://files.pythonhosted.org/packages/c7/28/2189571e711ab31cb7f682c739d05f033ae0e83e12eb4d5d21473ca2cd1a/wrapt-2.5.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d108bd98ab2c5ccbf55cf2154e0a73341df930030df2e01ee35d71a44a99d3ad", size = 250770, upload-time = "2026-10-29T23:04:49.183Z" }, + { url = "https://files.pythonhosted.org/packages/80/8e/46a6cc2e8215df7fae292178a355df729547ff9b912686823201d38be269/wrapt-2.5.2-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:858511731d6ede4217960c5fecdb5ecbf90e306dfd93b837395658d61374a498", size = 230031, upload-time = "2026-10-29T23:04:50.641Z" }, + { url = "https://files.pythonhosted.org/packages/97/38/dff531d0125fe880e05fe5fcb89f144d25ffbe95aa9a2c50539eabde7567/wrapt-2.5.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3b6885250b809532a634acd56e360d57c201c1cdc75cbf2b5e106241ece6cbc3", size = 247325, upload-time = "2026-10-29T23:04:52.16Z" }, + { url = "https://files.pythonhosted.org/packages/f0/5b/f4a78df4c52175d53ff095f2f5a54206a2d16e542b1dd2cdcaed901fe432/wrapt-2.5.2-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:c60a5cd8b7b510056f2643d94f1c144533a568ade2891a9dd35733f5e0055777", size = 226900, upload-time = "2026-10-29T23:04:53.559Z" }, + { url = "https://files.pythonhosted.org/packages/20/6f/8eab144c717300f76f1337ba58091f4de5da0f4d772727c9f41135b325c1/wrapt-2.5.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:40e03506f09dbf4854fd69848e843dadd9b7fbbdd217837cb27b37fb8f0b5bc3", size = 239017, upload-time = "2026-10-29T23:04:55.293Z" }, + { url = "https://files.pythonhosted.org/packages/59/1a/98e278ef2f8952a34f847a91f75899f6385762b07577c51a5ef82c768dfb/wrapt-2.5.2-cp312-cp312-win32.whl", hash = "sha256:ac00ba7d22d55d3c8d085005ce4942ff3bb118d480f889b5984f1d3454052917", size = 100660, upload-time = "2026-10-29T23:04:56.703Z" }, + { url = "https://files.pythonhosted.org/packages/7d/0a/ab645f8b4c594cac2db49bbc805bff97a9d3f2944061b865f67e0351b6c8/wrapt-2.5.2-cp312-cp312-win_amd64.whl", hash = "sha256:36021ef761856bdbc10d5f3da9a95af93ef60d6acf9c72c8f849937ba2b902aa", size = 106220, upload-time = "2026-10-29T23:04:58.081Z" }, + { url = "https://files.pythonhosted.org/packages/37/1e/232868b68b925ab29ca5acbec4ad54ee42a853fcd8f596bf7e0e0e3850f1/wrapt-2.5.2-cp312-cp312-win_arm64.whl", hash = "sha256:81bdf0bd141d9b1706a0751b2ee068c7afe6be4042712d427b32b3072d5c40b9", size = 103377, upload-time = "2026-10-29T23:04:59.366Z" }, + { url = "https://files.pythonhosted.org/packages/77/8b/1ea57b936b70f98e8e7620a352578096d778dbee95b262f73c066d210ffb/wrapt-2.5.2-py3-none-any.whl", hash = "sha256:61eb54e9540ea522c2c4752a57549a83d3c67e6704ed708403d28e067aaf90f0", size = 82025, upload-time = "2026-10-29T23:06:38.56Z" }, +] + +[[package]] +name = "zipp" +version = "4.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/18/05/ba1458dee0ac485b2a01071c1b88f5199fe878a485df482e10e1ef2205cc/zipp-4.1.2.tar.gz", hash = "sha256:cf474b90b07ea932637e17c6b9479de2b2882900b0d311bdaf4f5c000a15a97f", size = 27864, upload-time = "2026-10-21T18:57:04.743Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b8/44/7bd2007260da11cccca29ba2d4c345f1287aa663db38ee7cefeabf6386a5/zipp-4.1.2-py3-none-any.whl", hash = "sha256:b30cb6cd253122169c47de245ca24a86bac2a5d3802ba15bf7ba6414eb9e144d", size = 10581, upload-time = "2026-10-21T18:57:03.79Z" }, +]