diff --git a/.github/workflows/sim-tests.yml b/.github/workflows/sim-tests.yml new file mode 100644 index 0000000..6a2f840 --- /dev/null +++ b/.github/workflows/sim-tests.yml @@ -0,0 +1,58 @@ +name: Simulation tests + +on: + push: + pull_request: + +jobs: + test: + runs-on: ubuntu-latest + timeout-minutes: 20 + env: + MUJOCO_GL: osmesa + PYOPENGL_PLATFORM: osmesa + steps: + - uses: actions/checkout@v4 + - name: Verify checkpoint-free release contents + run: | + forbidden="$(git ls-files | grep -E '(^|/)(outputs|checkpoints|runs|wandb)/|\.(pt|pth|ckpt|onnx|h5|hdf5|pkl|pickle)$' || true)" + if [ -n "$forbidden" ]; then + echo "Generated model artifacts must not be tracked:" + echo "$forbidden" + exit 1 + fi + - uses: actions/setup-python@v5 + with: + python-version: "3.10" + cache: pip + - name: Install headless MuJoCo rendering libraries + run: | + sudo apt-get update + sudo apt-get install --yes libgl1 libosmesa6 + - name: Install public package and test dependencies + run: | + python -m pip install --upgrade pip + python -m pip install -e ".[rl,learned,test]" build twine + - name: Run tests + run: pytest -q + - name: Run supported integration commands + run: | + speedtuning-sim + speedtuning-sim --speed 1.5 + speedtuning-check-chunks + speedtuning-rainbow-poc + speedtuning-eval-speed --task tea_bag --base-policy recorded-chunk --episodes 1 + speedtuning-sweep --task tea_bag --base-policy scripted --speed-start 1.0 --speed-stop 1.1 --episodes-per-speed 1 --output /tmp/sweep.json + - name: Build source and wheel distributions + run: | + python -m build + python -m twine check dist/* + tar -tzf dist/*.tar.gz | grep 'docs/SCRIPTED_REPRODUCTION.md' + tar -tzf dist/*.tar.gz | grep 'benchmarks/scripted_results.json' + - name: Verify the wheel and packaged MuJoCo assets + run: | + python -m venv /tmp/speedtuning-package-check + /tmp/speedtuning-package-check/bin/python -m pip install dist/*.whl + cd /tmp + /tmp/speedtuning-package-check/bin/python -c "from experiment_config import load_experiment_config; assert load_experiment_config('scripted-pick-and-place')[0]['decisions'] == 100000" + /tmp/speedtuning-package-check/bin/speedtuning-sim --task tea_bag diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f39a8db --- /dev/null +++ b/.gitignore @@ -0,0 +1,35 @@ +# Python +__pycache__/ +*.py[cod] +.pytest_cache/ +.coverage +htmlcov/ +.mypy_cache/ +.ruff_cache/ + +# Environments and editors +.venv/ +.venv-legacy/ +.vscode/ +.idea/ +**/.DS_Store + +# Packaging +build/ +dist/ +*.egg-info/ + +# Generated research artifacts +outputs/ +results/ +logs/ +tmp/ +.tmp/ +data/ +data_local/ +checkpoints/ +*.pt +*.pth +*.ckpt +wandb/ +_wandb/ diff --git a/.python-version b/.python-version new file mode 100644 index 0000000..c8cfe39 --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.10 diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..96b7602 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,28 @@ +# Changelog + +## 0.1.0 + +- Released pick-and-place, insertion, and tea-bag MuJoCo tasks. +- Added parameterized execution speed for retained waypoint policies. +- Added a model-agnostic variable-speed action-chunk interface. +- Added external chunk-policy and speed-policy factory loading. +- Added supported Rainbow DQN speed-policy training, checkpoints, and evaluation. +- Added decision-level speed execution with fresh receding-horizon chunks and + shared frame-skip semantics for training and evaluation. +- Added stacked proprioceptive/visual speed observations with pretrained, random, + and external image-encoder support. +- Added retained ACT checkpoint/backbone adapters, checkpointed preprocessing, + seeded physical-acceleration metrics, fixed-speed sweeps, and plotting. +- Added an archival paper configuration and manifests for every published + simulation ablation. +- Added runnable scripted-policy presets for pick-and-place, insertion, and tea + bag, including the retained reward and Rainbow update schedule. +- Added from-scratch reproduction instructions and machine-readable reference + results for all three simulated tasks. +- Added seeded tea-bag pose randomization as a separately labeled robustness + protocol while preserving the fixed-pose historical environment. +- Added periodic training snapshots, task/protocol metadata validation, and safe + loading for locally generated speed-policy checkpoints. +- Added clean-install packaging, continuous integration, and release tests. +- Removed real-robot, private-path, scratch-output, and trained-checkpoint + artifacts from the public surface. diff --git a/CITATION.cff b/CITATION.cff new file mode 100644 index 0000000..4f7c297 --- /dev/null +++ b/CITATION.cff @@ -0,0 +1,25 @@ +cff-version: 1.2.0 +message: "If you use this software, please cite the SpeedTuning paper." +title: "SpeedTuning simulation and speed-policy infrastructure" +type: software +version: 0.1.0 +authors: + - family-names: Yuan + given-names: David D. +license: MIT +repository-code: "https://github.com/DaivdYuan/SpeedTuning" +preferred-citation: + type: conference-paper + title: "SpeedTuning: Speeding Up Policy Execution with Lightweight Reinforcement Learning" + authors: + - family-names: Yuan + given-names: David D. + - family-names: Zhao + given-names: Tony Z. + - family-names: Burns + given-names: Kaylee + - family-names: Finn + given-names: Chelsea + collection-title: "2025 IEEE International Conference on Robotics and Automation (ICRA)" + year: 2025 + doi: "10.1109/ICRA55743.2025.11128753" diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..6386907 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,21 @@ +# Contributing + +Bug reports and focused pull requests for the supported simulation surface are +welcome. Before opening a pull request: + +1. Install Python 3.10 and the development extras with + `uv sync --extra rl --extra learned --extra test`. +2. Run `uv run pytest -q`. +3. Run `uv run speedtuning-sim` and `uv run speedtuning-check-chunks` when + changing tasks, policies, interpolation, or physics assets. +4. Run `uv run speedtuning-rainbow-poc` when changing speed-policy learning. +5. Run a two-point `speedtuning-sweep` smoke test when changing metrics, + decision timing, or experiment manifests. + +Please keep real-robot dependencies, private checkpoints, datasets, machine-local +paths, and generated outputs outside this repository. New external policy support +should use the public adapters instead of adding a dependency on another research +repository to the core environment. + +By contributing, you agree that your contribution may be distributed under the +license applicable to the directory you modify. diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..686d48e --- /dev/null +++ b/LICENSE @@ -0,0 +1,22 @@ +MIT License + +Copyright (c) 2023 Tony Z. Zhao +Copyright (c) 2024-2026 David D. Yuan + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..cab13be --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,7 @@ +include CHANGELOG.md +include CITATION.cff +include CONTRIBUTING.md +include NOTICE.md +include requirements-sim.txt +recursive-include benchmarks *.json *.md +recursive-include docs *.md *.png diff --git a/NOTICE.md b/NOTICE.md new file mode 100644 index 0000000..5c5eace --- /dev/null +++ b/NOTICE.md @@ -0,0 +1,22 @@ +# Notices and attribution + +SpeedTuning simulation and speed-policy infrastructure includes code and assets +derived from [Action Chunking with Transformers +(ACT)](https://github.com/tonyzhaozh/act), originally released under the MIT +License. The original Tony Z. Zhao copyright notice is retained in `LICENSE`. + +The files under `detr/` are modified from +[DETR](https://github.com/facebookresearch/detr) and are distributed under the +Apache License 2.0 included at `detr/LICENSE`. + +The ALOHA/ViperX MuJoCo XML files and meshes under `assets/` were inherited from +the MIT-licensed ACT repository history. Task-specific tea-bag environment files +were added in the SpeedTuning development history and are distributed under this +repository's MIT License. + +SpeedTuning simulator recovery, public integration, and release engineering: +Copyright (c) 2024-2026 David D. Yuan. + +The README teaser image is rendered from the SpeedTuning project-page figure, +Copyright (c) the SpeedTuning authors and shared under CC BY-SA 4.0. The source +project page is https://daivdyuan.github.io/speed-tuning/. diff --git a/README.md b/README.md index fa59c36..ec65482 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,245 @@ -# SpeedTuning +
-This repository contains the code for *SpeedTuning: Speeding Up Policy Execution with Lightweight Reinforcement Learning* +# SpeedTuning: Speeding Up Policy Execution with Lightweight Reinforcement Learning -Note: This is a work in progress. The official release is expected by late September. +[David D. Yuan](https://www.linkedin.com/in/dewei-yuan) · +[Tony Z. Zhao](https://tonyzhaozh.github.io/) · +[Kaylee Burns](https://kayburns.github.io/) · +[Chelsea Finn](https://ai.stanford.edu/~cbfinn/) + +Stanford University · ICRA 2025 + +[Project Page](https://daivdyuan.github.io/speed-tuning/) · +[arXiv](https://arxiv.org/abs/2608.09138) · +[Conference Paper](https://daivdyuan.github.io/speed-tuning/static/pdfs/speedtuning_icra.pdf) · +[Video](https://daivdyuan.github.io/speed-tuning/static/videos/icra2025_final.mp4) · +[Simulation Reproduction](docs/SCRIPTED_REPRODUCTION.md) + +[![Tests](https://github.com/DaivdYuan/SpeedTuning/actions/workflows/sim-tests.yml/badge.svg)](https://github.com/DaivdYuan/SpeedTuning/actions/workflows/sim-tests.yml) + +
+ +

+ + SpeedTuning method and adaptive speed profile + +

+ +

+SpeedTuning keeps a base manipulation policy fixed and learns a lightweight +speed policy that accelerates safe phases while preserving precision around +critical interactions. +

+ +> [!NOTE] +> This repository is the simulation reproduction release. It provides complete +> from-scratch speed-policy training with bundled scripted task policies. + +## Overview + +Imitation-learned manipulation policies often inherit the operator's pace and +the hardware constraints present during data collection. Applying one global +interpolation factor can make execution faster, but it cannot distinguish +between transit phases that tolerate aggressive acceleration and contact-rich +phases that require precision. + +SpeedTuning adds a small reinforcement-learning policy on top of a frozen base +policy. At each decision, it selects a speed multiplier from the current robot +and task observation. The base policy continues to predict actions; SpeedTuning +only changes how quickly those actions are executed. + +This release supports the full simulation loop: + +1. run a task with a fixed scripted base policy; +2. train a Rainbow DQN policy over discrete speed multipliers; +3. evaluate success against physical acceleration; +4. compare the adaptive policy with matched fixed-speed baselines. + +## Included tasks + +| Task | Simulator objective | Public preset | +| --- | --- | --- | +| Pick-and-place | Transfer a cube between grippers | `scripted-pick-and-place` | +| Insertion | Insert a peg into a socket | `scripted-insertion` | +| Tea bag | Move a tea bag into a cup | `scripted-tea-bag` | + +An additional `scripted-tea-bag-randomized` preset samples initial tea-bag poses +for distributional evaluation. The retained fixed-pose environment remains +available for historical parity. + +## Installation + +Python 3.10 is required. MuJoCo and DM Control are pinned because contact +dynamics affect the scripted policies. + +Using `uv`: + +```bash +git clone https://github.com/DaivdYuan/SpeedTuning.git +cd SpeedTuning + +uv sync --extra test +uv run speedtuning-sim +``` + +Using `pip`: + +```bash +python3.10 -m venv .venv +source .venv/bin/activate +python -m pip install -e ".[test]" +speedtuning-sim +``` + +On a headless Linux machine, prefix simulator commands with `MUJOCO_GL=egl`. + +## Quick simulation check + +Run all three scripted tasks at nominal speed: + +```bash +uv run speedtuning-sim +``` + +Run one task with a fixed `1.5x` speed multiplier: + +```bash +uv run speedtuning-sim --task insertion --speed 1.5 +``` + +Each command prints a JSON summary and exits nonzero if the task fails. + +## Train a speed policy + +Install the reinforcement-learning extra and run a short CPU smoke test: + +```bash +uv sync --extra rl --extra test + +uv run speedtuning-train-speed \ + --config scripted-tea-bag \ + --task tea_bag \ + --decisions 1000 \ + --checkpoint-interval 0 \ + --output outputs/smoke_test.pt +``` + +For a full 100,000-decision run, use the preset without the smoke-test +overrides: + +```bash +uv run speedtuning-train-speed \ + --config scripted-tea-bag \ + --task tea_bag \ + --output outputs/tea_bag_speed.pt \ + --report outputs/tea_bag_speed.training.json +``` + +Training defaults to CPU. Add `--device cuda` when CUDA is available. Hardware +changes wall-clock time, not the simulation protocol or acceleration metric. + +Every full preset trains a separate task-specific policy. Generated checkpoints +and reports are written under the ignored `outputs/` directory; no pretrained +artifact is required or distributed. + +## Evaluate + +Evaluate the learned speed policy: + +```bash +uv run speedtuning-eval-speed \ + --config scripted-tea-bag \ + --task tea_bag \ + --speed-policy rainbow \ + --speed-checkpoint outputs/tea_bag_speed.pt \ + --episodes 20 +``` + +Measure a fixed-speed frontier: + +```bash +uv run speedtuning-sweep \ + --config scripted-tea-bag \ + --task tea_bag \ + --speed-start 1.0 --speed-stop 3.0 --speed-step 0.25 \ + --episodes-per-speed 20 \ + --output outputs/tea_bag_sweep.json +``` + +Physical acceleration is the nominal task horizon divided by the number of +executed MuJoCo steps. It is not the arithmetic mean of commanded multipliers. + +## Simulation reference results + +One seeded run using the final checkpoint from each 100,000-decision training +run produced: + +| Protocol | Adaptive SpeedTuning | Matched fixed speed | +| --- | --- | --- | +| Pick-and-place | 98% success at 3.856x | 66% at 3.846x | +| Insertion | 97% success at 2.387x | 52% at 2.381x | +| Tea bag, randomized poses | 78% success at 2.077x | 24% at 2.075x | + +These are reference points in the pinned simulator, not exact-decimal +guarantees. Reinforcement learning is stochastic; reruns should be compared by +the success/acceleration tradeoff. + +See the [full reproduction guide](docs/SCRIPTED_REPRODUCTION.md) for reward and +update definitions, all-task commands, pose protocols, and seeded evaluation. +The compact machine-readable record is +[`benchmarks/scripted_results.json`](benchmarks/scripted_results.json). + +## Bring your own task policy + +The speed controller can wrap an external policy that returns action chunks with +shape `[time, 14]`. A `module:factory` adapter makes it possible to train or +evaluate another repository's task policy without modifying this codebase. + +See [External task-policy integration](docs/EXTERNAL_POLICIES.md) for: + +- the Python and CLI interfaces; +- ACT checkpoint and normalization support; +- visual, state, and external speed-policy observations; +- the archival learned-policy configuration and ablations. + +## Commands + +| Command | Purpose | +| --- | --- | +| `speedtuning-sim` | Run scripted simulator tasks at a fixed speed | +| `speedtuning-train-speed` | Train a Rainbow speed policy | +| `speedtuning-eval-speed` | Evaluate fixed, profiled, or learned speed policies | +| `speedtuning-sweep` | Build a success-versus-acceleration curve | +| `speedtuning-check-chunks` | Validate action-chunk integration | +| `speedtuning-rainbow-poc` | Run a small Rainbow optimization check | + +## Scope and limitations + +- This release reproduces the methodology with scripted base policies in + simulation; it does not claim to reproduce the paper's learned-ACT table. +- Real-robot execution is not part of the supported API. +- External task policies remain responsible for their architectures, + preprocessing, normalization statistics, and checkpoint compatibility. +- The physics stack is intentionally pinned for reproducibility. + +## Citation + +If you use this code, please cite: + +```bibtex +@inproceedings{yuan2025speedtuning, + title = {{SpeedTuning}: Speeding Up Policy Execution with Lightweight Reinforcement Learning}, + author = {Yuan, David D. and Zhao, Tony Z. and Burns, Kaylee and Finn, Chelsea}, + booktitle = {2025 IEEE International Conference on Robotics and Automation (ICRA)}, + year = {2025}, + doi = {10.1109/ICRA55743.2025.11128753} +} +``` + +Citation metadata is also available in [`CITATION.cff`](CITATION.cff). + +## License and acknowledgments + +SpeedTuning is released under the MIT License. The ACT-derived DETR code under +`detr/` retains its Apache-2.0 license. Simulator assets and upstream attribution +are documented in [`NOTICE.md`](NOTICE.md). diff --git a/act_integration.py b/act_integration.py new file mode 100644 index 0000000..9ac89c0 --- /dev/null +++ b/act_integration.py @@ -0,0 +1,242 @@ +"""Load retained ACT checkpoints through the public chunk-policy interface.""" + +from __future__ import annotations + +import json +import pickle +from pathlib import Path + +import numpy as np + +from chunked_policy import TorchChunkPredictor + + +REQUIRED_STATS = ("qpos_mean", "qpos_std", "action_mean", "action_std") + + +def _load_mapping(path): + path = Path(path) + if path.suffix == ".npz": + with np.load(path) as values: + return {key: values[key] for key in values.files} + if path.suffix == ".json": + return json.loads(path.read_text()) + if path.suffix in {".pkl", ".pickle"}: + with path.open("rb") as stream: + return pickle.load(stream) + raise ValueError("ACT stats must use .npz, .json, .pkl, or .pickle") + + +def _checkpoint_parts(checkpoint, device): + try: + import torch + except ImportError as exc: + raise RuntimeError("ACT integration requires: uv sync --extra learned") from exc + if checkpoint is None: + raise ValueError("An ACT checkpoint path is required") + # ACT checkpoints may include NumPy normalization arrays and legacy config + # objects. They therefore require pickle loading and must come from a + # trusted source (normally the user's own task-policy training run). + payload = torch.load( + Path(checkpoint), map_location=device, weights_only=False + ) + if not isinstance(payload, dict): + raise ValueError("ACT checkpoint must contain a state dictionary or payload") + for key in ("model_state_dict", "policy_state_dict", "state_dict"): + if key in payload: + return payload, payload[key] + # A raw torch state dictionary maps names to tensors. + if payload and all(isinstance(key, str) for key in payload): + return {}, payload + raise ValueError("ACT checkpoint does not contain recognizable model weights") + + +def _resolve_config(payload, policy_config, camera_names, device): + config = dict(payload.get("policy_config", {})) + config.update(policy_config or {}) + if camera_names is not None: + config["camera_names"] = list(camera_names) + if not config.get("camera_names"): + raise ValueError("ACT policy_config must provide camera_names") + required = ("num_queries", "hidden_dim", "dim_feedforward", "enc_layers", "dec_layers", "nheads") + missing = [key for key in required if key not in config] + if missing: + raise ValueError(f"ACT policy_config is missing: {', '.join(missing)}") + config.setdefault("lr", 1e-4) + config.setdefault("lr_backbone", 0.0) + config.setdefault("kl_weight", 10.0) + config.setdefault("backbone", "resnet18") + config.setdefault("pretrained_backbone", False) + config["device"] = device + return config + + +def _resolve_stats(payload, stats_path): + stats = dict(payload.get("stats", {})) + if stats_path is not None: + stats.update(_load_mapping(stats_path)) + missing = [key for key in REQUIRED_STATS if key not in stats] + if missing: + raise ValueError( + "ACT normalization stats are missing: " + ", ".join(missing) + ) + return {key: np.asarray(stats[key], dtype=np.float32) for key in REQUIRED_STATS} + + +def load_act_policy( + checkpoint, + device="cpu", + stats_path=None, + policy_config=None, + camera_names=None, + strict=True, +): + """Return the ACT module, resolved configuration, and normalization stats.""" + + from policy import ACTPolicy + + payload, state_dict = _checkpoint_parts(checkpoint, device) + config = _resolve_config(payload, policy_config, camera_names, device) + stats = _resolve_stats(payload, stats_path) + model = ACTPolicy(config) + incompatible = model.load_state_dict(state_dict, strict=bool(strict)) + if not strict and (incompatible.missing_keys or incompatible.unexpected_keys): + # Keep the information available to callers without printing during imports. + model.checkpoint_incompatibilities = { + "missing_keys": list(incompatible.missing_keys), + "unexpected_keys": list(incompatible.unexpected_keys), + } + model.eval() + return model, config, stats + + +def build_act_chunk_predictor( + task_name, + checkpoint, + device="cpu", + stats_path=None, + policy_config=None, + camera_names=None, + strict=True, +): + """Factory usable as ``act_integration:build_act_chunk_predictor``.""" + + del task_name + model, config, stats = load_act_policy( + checkpoint=checkpoint, + device=device, + stats_path=stats_path, + policy_config=policy_config, + camera_names=camera_names, + strict=strict, + ) + return TorchChunkPredictor( + model=model, + camera_names=config["camera_names"], + qpos_mean=stats["qpos_mean"], + qpos_std=stats["qpos_std"], + action_mean=stats["action_mean"], + action_std=stats["action_std"], + device=device, + ) + + +class ACTBackboneObservationEncoder: + """Use a supplied ACT task policy's ResNet backbone for speed features.""" + + requires_images = True + + def __init__(self, model, camera_names, include_qvel=True, device="cpu"): + import torch + + self.torch = torch + self.device = torch.device(device) + self.camera_names = tuple(camera_names) + self.include_qvel = bool(include_qvel) + self.backbone = model.model.backbones[0] + self.backbone.to(self.device).eval() + self.feature_dim = int(self.backbone.num_channels) + self.mean = torch.tensor( + [0.485, 0.456, 0.406], dtype=torch.float32, device=self.device + ).view(1, 3, 1, 1) + self.std = torch.tensor( + [0.229, 0.224, 0.225], dtype=torch.float32, device=self.device + ).view(1, 3, 1, 1) + + def reset(self): + return None + + def __call__(self, observation): + torch = self.torch + if "images" not in observation: + raise ValueError("ACT-backbone speed observations require images") + images = np.stack( + [observation["images"][name] for name in self.camera_names] + ).transpose(0, 3, 1, 2) + tensor = torch.as_tensor(images, dtype=torch.float32, device=self.device) / 255.0 + tensor = (tensor - self.mean) / self.std + features = [] + with torch.inference_mode(): + for image in tensor: + backbone_features, _ = self.backbone(image.unsqueeze(0)) + feature_map = backbone_features[-1] + features.append(feature_map.mean(dim=(2, 3)).squeeze(0)) + proprioception = [np.asarray(observation["qpos"], dtype=np.float32)] + if self.include_qvel: + proprioception.append(np.asarray(observation["qvel"], dtype=np.float32)) + return np.concatenate( + proprioception + + [torch.cat(features).detach().cpu().numpy().astype(np.float32)] + ) + + def output_dim(self, env_state_dim): + del env_state_dim + return 14 + (14 if self.include_qvel else 0) + len(self.camera_names) * self.feature_dim + + def spec(self): + return { + "type": "act_backbone", + "camera_names": list(self.camera_names), + "include_qpos": True, + "include_qvel": self.include_qvel, + "include_env_state": False, + "feature_dim": self.feature_dim, + } + + def state_dict(self): + return { + key: value.detach().cpu() + for key, value in self.backbone.state_dict().items() + } + + def load_state_dict(self, state_dict): + self.backbone.load_state_dict(state_dict) + + +def build_act_observation_encoder( + task_name, + checkpoint, + device="cpu", + stats_path=None, + policy_config=None, + camera_names=None, + include_qvel=True, + strict=True, +): + """Factory for the task-policy image-encoder ablation.""" + + del task_name + model, config, _ = load_act_policy( + checkpoint=checkpoint, + device=device, + stats_path=stats_path, + policy_config=policy_config, + camera_names=camera_names, + strict=strict, + ) + return ACTBackboneObservationEncoder( + model, + config["camera_names"], + include_qvel=include_qvel, + device=device, + ) diff --git a/assets/__init__.py b/assets/__init__.py new file mode 100644 index 0000000..9d30322 --- /dev/null +++ b/assets/__init__.py @@ -0,0 +1 @@ +"""Packaged MuJoCo models and meshes for the simulation tasks.""" diff --git a/assets/bimanual_viperx_ee_insertion.xml b/assets/bimanual_viperx_ee_insertion.xml new file mode 100644 index 0000000..700aaac --- /dev/null +++ b/assets/bimanual_viperx_ee_insertion.xml @@ -0,0 +1,59 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/assets/bimanual_viperx_ee_transfer_cube.xml b/assets/bimanual_viperx_ee_transfer_cube.xml new file mode 100644 index 0000000..2589384 --- /dev/null +++ b/assets/bimanual_viperx_ee_transfer_cube.xml @@ -0,0 +1,48 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/assets/bimanual_viperx_ee_transfer_tea_bag.xml b/assets/bimanual_viperx_ee_transfer_tea_bag.xml new file mode 100644 index 0000000..3358ebc --- /dev/null +++ b/assets/bimanual_viperx_ee_transfer_tea_bag.xml @@ -0,0 +1,103 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/assets/bimanual_viperx_insertion.xml b/assets/bimanual_viperx_insertion.xml new file mode 100644 index 0000000..f701d70 --- /dev/null +++ b/assets/bimanual_viperx_insertion.xml @@ -0,0 +1,53 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/assets/bimanual_viperx_transfer_cube.xml b/assets/bimanual_viperx_transfer_cube.xml new file mode 100644 index 0000000..bdc9e64 --- /dev/null +++ b/assets/bimanual_viperx_transfer_cube.xml @@ -0,0 +1,42 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/assets/bimanual_viperx_transfer_tea_bag.xml b/assets/bimanual_viperx_transfer_tea_bag.xml new file mode 100644 index 0000000..51974c8 --- /dev/null +++ b/assets/bimanual_viperx_transfer_tea_bag.xml @@ -0,0 +1,97 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/assets/scene.xml b/assets/scene.xml new file mode 100644 index 0000000..ae59450 --- /dev/null +++ b/assets/scene.xml @@ -0,0 +1,38 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/assets/tabletop.stl b/assets/tabletop.stl new file mode 100644 index 0000000..ab35cdf Binary files /dev/null and b/assets/tabletop.stl differ diff --git a/assets/vx300s_10_custom_finger_left.stl b/assets/vx300s_10_custom_finger_left.stl new file mode 100644 index 0000000..534c7af Binary files /dev/null and b/assets/vx300s_10_custom_finger_left.stl differ diff --git a/assets/vx300s_10_custom_finger_right.stl b/assets/vx300s_10_custom_finger_right.stl new file mode 100644 index 0000000..d6a492c Binary files /dev/null and b/assets/vx300s_10_custom_finger_right.stl differ diff --git a/assets/vx300s_1_base.stl b/assets/vx300s_1_base.stl new file mode 100644 index 0000000..5a7efda Binary files /dev/null and b/assets/vx300s_1_base.stl differ diff --git a/assets/vx300s_2_shoulder.stl b/assets/vx300s_2_shoulder.stl new file mode 100644 index 0000000..dc22aa7 Binary files /dev/null and b/assets/vx300s_2_shoulder.stl differ diff --git a/assets/vx300s_3_upper_arm.stl b/assets/vx300s_3_upper_arm.stl new file mode 100644 index 0000000..111c586 Binary files /dev/null and b/assets/vx300s_3_upper_arm.stl differ diff --git a/assets/vx300s_4_upper_forearm.stl b/assets/vx300s_4_upper_forearm.stl new file mode 100644 index 0000000..8170d21 Binary files /dev/null and b/assets/vx300s_4_upper_forearm.stl differ diff --git a/assets/vx300s_5_lower_forearm.stl b/assets/vx300s_5_lower_forearm.stl new file mode 100644 index 0000000..39581f8 Binary files /dev/null and b/assets/vx300s_5_lower_forearm.stl differ diff --git a/assets/vx300s_6_wrist.stl b/assets/vx300s_6_wrist.stl new file mode 100644 index 0000000..ab8423e Binary files /dev/null and b/assets/vx300s_6_wrist.stl differ diff --git a/assets/vx300s_7_gripper.stl b/assets/vx300s_7_gripper.stl new file mode 100644 index 0000000..043db9c Binary files /dev/null and b/assets/vx300s_7_gripper.stl differ diff --git a/assets/vx300s_8_gripper_prop.stl b/assets/vx300s_8_gripper_prop.stl new file mode 100644 index 0000000..36099b4 Binary files /dev/null and b/assets/vx300s_8_gripper_prop.stl differ diff --git a/assets/vx300s_9_gripper_bar.stl b/assets/vx300s_9_gripper_bar.stl new file mode 100644 index 0000000..eba3caa Binary files /dev/null and b/assets/vx300s_9_gripper_bar.stl differ diff --git a/assets/vx300s_dependencies.xml b/assets/vx300s_dependencies.xml new file mode 100644 index 0000000..c75d3ad --- /dev/null +++ b/assets/vx300s_dependencies.xml @@ -0,0 +1,17 @@ + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/assets/vx300s_left.xml b/assets/vx300s_left.xml new file mode 100644 index 0000000..61e6219 --- /dev/null +++ b/assets/vx300s_left.xml @@ -0,0 +1,59 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/assets/vx300s_right.xml b/assets/vx300s_right.xml new file mode 100644 index 0000000..2c6f007 --- /dev/null +++ b/assets/vx300s_right.xml @@ -0,0 +1,59 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 0000000..5dbccc8 --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,18 @@ +# Reference results + +[`scripted_results.json`](scripted_results.json) records one seeded run of the +from-scratch scripted-policy protocol. The record is hardware-neutral: it +reports simulator success and physical acceleration, not training time. + +| Protocol | Learned speed | Matched fixed speed | +| --- | --- | --- | +| Pick-and-place | 98% at 3.856x | 66% at 3.846x | +| Insertion | 97% at 2.387x | 52% at 2.381x | +| Tea bag, randomized poses | 78% at 2.077x | 24% at 2.075x | + +The JSON includes the preset, training seed, and episode count needed to +interpret each result. Training is stochastic, so these values +are reference points rather than exact-decimal guarantees. + +See the [reproduction guide](../docs/SCRIPTED_REPRODUCTION.md) for commands and +protocol details. No trained checkpoint is distributed with the repository. diff --git a/benchmarks/scripted_results.json b/benchmarks/scripted_results.json new file mode 100644 index 0000000..02610bc --- /dev/null +++ b/benchmarks/scripted_results.json @@ -0,0 +1,59 @@ +{ + "schema_version": 1, + "description": "Reference results for from-scratch scripted-policy training", + "protocol": { + "base_policy": "bundled scripted policy", + "training_decisions": 100000, + "evaluation_checkpoint": "final training output", + "held_out_seeds": [100, 199], + "episodes_per_task": 100, + "metric": "physical_acceleration=nominal_horizon/executed_physics_steps" + }, + "results": { + "pick_and_place": { + "preset": "scripted-pick-and-place", + "training_seed": 0, + "learned_speed": { + "successes": 98, + "success_rate": 0.98, + "mean_physical_acceleration": 3.856 + }, + "matched_fixed_speed": { + "commanded_speed": 3.856, + "successes": 66, + "success_rate": 0.66, + "mean_physical_acceleration": 3.846 + } + }, + "insertion": { + "preset": "scripted-insertion", + "training_seed": 0, + "learned_speed": { + "successes": 97, + "success_rate": 0.97, + "mean_physical_acceleration": 2.387 + }, + "matched_fixed_speed": { + "commanded_speed": 2.387, + "successes": 52, + "success_rate": 0.52, + "mean_physical_acceleration": 2.381 + } + }, + "tea_bag": { + "preset": "scripted-tea-bag-randomized", + "training_seed": 1, + "learned_speed": { + "successes": 78, + "success_rate": 0.78, + "mean_physical_acceleration": 2.077 + }, + "matched_fixed_speed": { + "commanded_speed": 2.075, + "successes": 24, + "success_rate": 0.24, + "mean_physical_acceleration": 2.075 + } + } + } +} diff --git a/chunked_policy.py b/chunked_policy.py new file mode 100644 index 0000000..e22efad --- /dev/null +++ b/chunked_policy.py @@ -0,0 +1,324 @@ +"""Adapters for upstream policies that predict joint-action chunks. + +The public contract is intentionally model-agnostic: a predictor receives one +DM Control observation dictionary and returns one or more 14D joint actions. +``ChunkPredictorAdapter`` handles common NumPy, PyTorch, tuple, and dictionary +outputs, while ``ChunkedPolicyRunner`` turns those chunks into individual +simulator actions at a parameterized execution speed. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable + +import numpy as np + +from constants import PUPPET_GRIPPER_POSITION_NORMALIZE_FN +from ee_sim_env import make_ee_sim_env +from scripted_policy import make_scripted_policy +from sim_env import make_sim_env +from sim_tasks import get_task_spec, normalize_task_name + + +def _to_numpy(value: Any) -> np.ndarray: + """Convert common model outputs without importing a framework eagerly.""" + + detach = getattr(value, "detach", None) + if detach is not None: + value = detach() + cpu = getattr(value, "cpu", None) + if cpu is not None: + value = cpu() + numpy = getattr(value, "numpy", None) + if numpy is not None: + value = numpy() + return np.asarray(value, dtype=np.float64) + + +def as_action_chunk(output: Any) -> np.ndarray: + """Normalize an upstream policy output to a finite ``[time, 14]`` array.""" + + if isinstance(output, dict): + for key in ("actions", "action", "chunk"): + if key in output: + output = output[key] + break + else: + raise ValueError( + "Chunk dictionaries must contain an 'actions', 'action', or 'chunk' key" + ) + if isinstance(output, (tuple, list)): + if not output: + raise ValueError("A chunked policy returned an empty sequence") + output = output[0] + + actions = _to_numpy(output) + if actions.ndim == 3 and actions.shape[0] == 1: + actions = actions[0] + if actions.ndim == 1 and actions.shape[0] == 14: + actions = actions[None] + if actions.ndim != 2 or actions.shape[1] != 14 or len(actions) == 0: + raise ValueError("Chunked policy output must have shape [time, 14]") + if not np.all(np.isfinite(actions)): + raise ValueError("Chunked policy produced a non-finite action") + return actions + + +class ChunkPredictorAdapter: + """Adapt an arbitrary upstream chunk predictor to the public contract. + + ``predictor`` can be callable or expose ``predict_chunk(observation)``. + Optional adapters make it possible to translate repository-specific + observations and outputs without changing the simulator integration. + """ + + def __init__( + self, + predictor: Any, + observation_adapter: Callable[[dict], Any] | None = None, + output_adapter: Callable[[Any], Any] | None = None, + ): + predict = getattr(predictor, "predict_chunk", None) + if predict is None and callable(predictor): + predict = predictor + if predict is None: + raise TypeError("predictor must be callable or define predict_chunk()") + self.predictor = predictor + self._predict = predict + self.observation_adapter = observation_adapter or (lambda observation: observation) + self.output_adapter = output_adapter or (lambda output: output) + + def reset(self): + reset = getattr(self.predictor, "reset", None) + if reset is not None: + reset() + + def advance(self, nominal_steps): + """Notify predictors that explicitly track nominal demonstration time.""" + + advance = getattr(self.predictor, "advance", None) + if advance is not None: + advance(float(nominal_steps)) + + def __call__(self, observation): + model_input = self.observation_adapter(observation) + output = self.output_adapter(self._predict(model_input)) + return as_action_chunk(output) + + +def interpolate_action_chunk(actions, speed=1.0): + """Resample a ``[time, 14]`` action chunk at a new execution speed.""" + + actions = as_action_chunk(actions) + if speed <= 0: + raise ValueError("speed must be positive") + sample_times = np.arange(0.0, len(actions), speed) + lower = np.floor(sample_times).astype(int) + upper = np.minimum(lower + 1, len(actions) - 1) + fraction = (sample_times - lower)[:, None] + return actions[lower] + fraction * (actions[upper] - actions[lower]) + + +class ChunkedPolicyRunner: + """Turn a chunk-predicting policy into one joint action per simulator step. + + Speed is expressed in nominal policy timesteps per physics step. It may be + fixed at construction or supplied for every call, allowing a learned speed + policy to change acceleration online without modifying the upstream model. + """ + + def __init__(self, predictor, speed=1.0): + self.predictor = ( + predictor + if isinstance(predictor, ChunkPredictorAdapter) + else ChunkPredictorAdapter(predictor) + ) + self.speed = self._validate_speed(speed) + self._chunk = None + self._chunk_index = 0.0 + + @staticmethod + def _validate_speed(speed): + speed = float(speed) + if not np.isfinite(speed) or speed <= 0: + raise ValueError("speed must be a finite positive value") + return speed + + def reset(self): + self._chunk = None + self._chunk_index = 0.0 + reset = getattr(self.predictor, "reset", None) + if reset is not None: + reset() + + def begin_decision(self, observation, speed=None): + """Predict a fresh receding-horizon chunk for one speed decision.""" + + speed = self.speed if speed is None else self._validate_speed(speed) + self._chunk = self.predictor(observation) + self._chunk_index = 0.0 + self._decision_speed = speed + return self._chunk + + def action(self, observation, speed=None): + speed = self.speed if speed is None else self._validate_speed(speed) + if self._chunk is None or self._chunk_index >= len(self._chunk): + self.begin_decision(observation, speed=speed) + + lower = int(np.floor(self._chunk_index)) + upper = min(lower + 1, len(self._chunk) - 1) + fraction = self._chunk_index - lower + action = self._chunk[lower] + fraction * ( + self._chunk[upper] - self._chunk[lower] + ) + self._chunk_index += speed + self.predictor.advance(speed) + return action.copy() + + +class TorchChunkPredictor: + """Pre/post-processing adapter for ACT-compatible PyTorch policies. + + The wrapped model must accept normalized ``qpos`` with shape ``[1, 14]`` and + RGB images with shape ``[1, cameras, 3, height, width]``, and return an action + tensor with shape ``[1, chunk, 14]``. + """ + + def __init__( + self, + model, + camera_names, + qpos_mean, + qpos_std, + action_mean, + action_std, + device=None, + ): + try: + import torch + except ImportError as exc: + raise RuntimeError("Learned policies require: uv sync --extra learned") from exc + self.torch = torch + self.model = model + self.camera_names = tuple(camera_names) + self.device = torch.device( + device or ("cuda" if torch.cuda.is_available() else "cpu") + ) + self.qpos_mean = np.asarray(qpos_mean) + self.qpos_std = np.maximum(np.asarray(qpos_std), 1e-6) + self.action_mean = np.asarray(action_mean) + self.action_std = np.asarray(action_std) + for value, name in ( + (self.qpos_mean, "qpos_mean"), + (self.qpos_std, "qpos_std"), + (self.action_mean, "action_mean"), + (self.action_std, "action_std"), + ): + if value.shape != (14,): + raise ValueError(f"{name} must have shape (14,)") + self.model.to(self.device) + self.model.eval() + + def __call__(self, observation): + torch = self.torch + if "images" not in observation: + raise ValueError("The environment must be created with render_images=True") + qpos = (np.asarray(observation["qpos"]) - self.qpos_mean) / self.qpos_std + images = np.stack( + [observation["images"][name] for name in self.camera_names], axis=0 + ).transpose(0, 3, 1, 2) + qpos_tensor = torch.as_tensor(qpos, dtype=torch.float32, device=self.device)[None] + image_tensor = torch.as_tensor( + images / 255.0, dtype=torch.float32, device=self.device + )[None] + with torch.inference_mode(): + output = self.model(qpos_tensor, image_tensor) + chunk = as_action_chunk(output) + return chunk * self.action_std + self.action_mean + + +@dataclass +class JointDemonstration: + actions: np.ndarray + object_pose: np.ndarray + + +def collect_scripted_joint_demonstration(task_name, seed=0): + """Convert the retained EE scripted rollout into joint commands.""" + + task_name = normalize_task_name(task_name) + spec = get_task_spec(task_name) + env = make_ee_sim_env(task_name, render_images=False, seed=seed) + timestep = env.reset() + episode = [timestep] + policy = make_scripted_policy(task_name) + for _ in range(spec.episode_len): + timestep = env.step(policy(timestep)) + episode.append(timestep) + + actions = [] + for item in episode: + action = item.observation["qpos"].copy() + gripper_ctrl = item.observation["gripper_ctrl"] + action[6] = PUPPET_GRIPPER_POSITION_NORMALIZE_FN(gripper_ctrl[0]) + action[13] = PUPPET_GRIPPER_POSITION_NORMALIZE_FN(gripper_ctrl[2]) + actions.append(action) + return JointDemonstration( + actions=np.asarray(actions), + object_pose=episode[0].observation["env_state"].copy(), + ) + + +class RecordedChunkPredictor: + """Deterministic chunk oracle used to validate the learned-policy contract.""" + + def __init__(self, actions, chunk_size): + self.actions = np.asarray(actions) + self.chunk_size = chunk_size + self.cursor = 0.0 + + def reset(self): + self.cursor = 0.0 + + def advance(self, nominal_steps): + self.cursor = min(self.cursor + float(nominal_steps), len(self.actions) - 1) + + def __call__(self, observation): + del observation + start = int(np.floor(self.cursor)) + end = min(start + self.chunk_size, len(self.actions)) + chunk = self.actions[start:end] + if len(chunk) < self.chunk_size: + chunk = np.concatenate( + [chunk, np.repeat(chunk[-1:], self.chunk_size - len(chunk), axis=0)] + ) + return chunk + + +def replay_recorded_chunks(task_name, chunk_size=25, seed=0): + """Validate chunked-policy replay through the joint-control environment.""" + + task_name = normalize_task_name(task_name) + demonstration = collect_scripted_joint_demonstration(task_name, seed=seed) + env = make_sim_env( + task_name, + render_images=False, + seed=seed, + object_pose=demonstration.object_pose, + ) + timestep = env.reset() + predictor = RecordedChunkPredictor(demonstration.actions, chunk_size) + runner = ChunkedPolicyRunner(predictor) + rewards = [] + for _ in range(len(demonstration.actions)): + timestep = env.step(runner.action(timestep.observation)) + rewards.append(int(timestep.reward or 0)) + return { + "task": task_name, + "success": max(rewards, default=0) == env.task.max_reward, + "max_reward": max(rewards, default=0), + "target_reward": env.task.max_reward, + "chunk_size": chunk_size, + "steps": len(demonstration.actions), + } diff --git a/configs/README.md b/configs/README.md new file mode 100644 index 0000000..bb1f71c --- /dev/null +++ b/configs/README.md @@ -0,0 +1,39 @@ +# Experiment manifests + +`paper_sim.json` is the best simulation configuration recoverable from the +published paper and the retained `SpeedTuningViz` experiment directory names. It +is an archival preset, not a guarantee of exact numerical reproduction. + +`scripted_tea_bag.json` is the fully runnable public reproduction. It uses the +included tea-bag waypoint policy and the latest retained training recipe from +the repository history: five speed actions, state observations, quadratic speed +reward, 100k decisions, and episode-boundary Rainbow updates. Load it with +`--config scripted-tea-bag`. + +`scripted_tea_bag_randomized.json` inherits that recipe and samples the tea bag +within the original ACT cube-position range at every reset. This opt-in variant +is the appropriate preset for reporting success rates across simulator seeds; +the exact historical tea-bag environment uses one fixed initial pose. + +`scripted_pick_and_place.json` and `scripted_insertion.json` apply the same +scripted-base Rainbow recipe to the other reconstructed tasks. Pick-and-place +includes speed actions through 4.5x because its transport phase tolerates the +paper's higher acceleration range; insertion retains the denser historical +1.0x-3.0x actions for contact-sensitive control. + +The main recovered choices are: + +- discrete speeds: 1.5, 2.0, 3.0, and 4.5; +- frame skip 10 and observation stack 5; +- proprioception plus top-camera features from pretrained ResNet-18; +- no privileged simulator object state; +- quadratic speed reward and hidden dimension 1024. + +The exact original task-policy checkpoints, speed-policy checkpoints, datasets, +and full training command were not recovered. Values retained only in the old +trainer, such as gamma and replay capacity, are included to make the preset +concrete and are identified by the manifest's archival status. + +Files in `ablations/` inherit the paper preset and override one experimental +factor. They can be passed directly to `--config`; command-line arguments may +override any value. diff --git a/configs/__init__.py b/configs/__init__.py new file mode 100644 index 0000000..042db17 --- /dev/null +++ b/configs/__init__.py @@ -0,0 +1 @@ +"""Packaged SpeedTuning experiment manifests.""" diff --git a/configs/ablations/act_encoder.json b/configs/ablations/act_encoder.json new file mode 100644 index 0000000..2a8aef3 --- /dev/null +++ b/configs/ablations/act_encoder.json @@ -0,0 +1,8 @@ +{ + "name": "act-task-policy-encoder", + "inherits": "paper-sim", + "observation": { + "speed_observation": "external", + "observation_encoder_loader": "act_integration:build_act_observation_encoder" + } +} diff --git a/configs/ablations/frame_skip_25.json b/configs/ablations/frame_skip_25.json new file mode 100644 index 0000000..0668c91 --- /dev/null +++ b/configs/ablations/frame_skip_25.json @@ -0,0 +1,5 @@ +{ + "name": "frame-skip-25", + "inherits": "paper-sim", + "environment": {"frame_skip": 25} +} diff --git a/configs/ablations/frame_skip_5.json b/configs/ablations/frame_skip_5.json new file mode 100644 index 0000000..fcdab6b --- /dev/null +++ b/configs/ablations/frame_skip_5.json @@ -0,0 +1,5 @@ +{ + "name": "frame-skip-5", + "inherits": "paper-sim", + "environment": {"frame_skip": 5} +} diff --git a/configs/ablations/no_image.json b/configs/ablations/no_image.json new file mode 100644 index 0000000..3402632 --- /dev/null +++ b/configs/ablations/no_image.json @@ -0,0 +1,11 @@ +{ + "name": "no-image", + "inherits": "paper-sim", + "observation": { + "speed_observation": "state", + "include_env_state": false, + "include_qpos": true, + "include_qvel": true, + "frame_stack": 5 + } +} diff --git a/configs/ablations/random_resnet.json b/configs/ablations/random_resnet.json new file mode 100644 index 0000000..5056f4b --- /dev/null +++ b/configs/ablations/random_resnet.json @@ -0,0 +1,5 @@ +{ + "name": "random-resnet18", + "inherits": "paper-sim", + "observation": {"image_encoder": "resnet18-random"} +} diff --git a/configs/ablations/reward_beta_0.json b/configs/ablations/reward_beta_0.json new file mode 100644 index 0000000..d2f31fe --- /dev/null +++ b/configs/ablations/reward_beta_0.json @@ -0,0 +1,5 @@ +{ + "name": "reward-beta-0", + "inherits": "paper-sim", + "reward": {"speed_power": 0.0, "speed_weight": 0.0} +} diff --git a/configs/ablations/reward_beta_1.json b/configs/ablations/reward_beta_1.json new file mode 100644 index 0000000..7a0ceef --- /dev/null +++ b/configs/ablations/reward_beta_1.json @@ -0,0 +1,5 @@ +{ + "name": "reward-beta-1", + "inherits": "paper-sim", + "reward": {"speed_power": 1.0} +} diff --git a/configs/ablations/reward_beta_3.json b/configs/ablations/reward_beta_3.json new file mode 100644 index 0000000..fc3bf4d --- /dev/null +++ b/configs/ablations/reward_beta_3.json @@ -0,0 +1,5 @@ +{ + "name": "reward-beta-3", + "inherits": "paper-sim", + "reward": {"speed_power": 3.0} +} diff --git a/configs/ablations/scripted_policy.json b/configs/ablations/scripted_policy.json new file mode 100644 index 0000000..c7d84b7 --- /dev/null +++ b/configs/ablations/scripted_policy.json @@ -0,0 +1,5 @@ +{ + "name": "scripted-policy", + "inherits": "paper-sim", + "base_policy": {"base_policy": "scripted"} +} diff --git a/configs/paper_sim.json b/configs/paper_sim.json new file mode 100644 index 0000000..a1be8dd --- /dev/null +++ b/configs/paper_sim.json @@ -0,0 +1,49 @@ +{ + "name": "paper-sim", + "status": "best-recovered archival configuration; checkpoints and data are not included", + "evidence": "Published paper plus retained SpeedTuningViz experiment names", + "base_policy": { + "base_policy": "external-chunk", + "chunk_size": 100 + }, + "environment": { + "speed_values": [1.5, 2.0, 3.0, 4.5], + "frame_skip": 10 + }, + "observation": { + "speed_observation": "visual", + "camera_names": ["top"], + "image_encoder": "resnet18-pretrained", + "image_size": 224, + "include_env_state": false, + "include_qpos": true, + "include_qvel": true, + "frame_stack": 5 + }, + "reward": { + "success_bonus": 100.0, + "speed_weight": 0.01, + "speed_power": 2.0 + }, + "training": { + "decisions": 100000, + "memory_size": 1000000, + "batch_size": 128, + "learning_starts": 512, + "gradient_steps": 1, + "update_schedule": "episode", + "checkpoint_interval": 10000, + "hidden_dim": 1024, + "gamma": 0.99, + "target_update": 50, + "beta_schedule": "legacy" + }, + "evaluation": { + "episodes": 2000, + "adaptive_episodes": 2000, + "episodes_per_speed": 100, + "speed_start": 1.0, + "speed_stop": 4.5, + "speed_step": 0.1 + } +} diff --git a/configs/scripted_insertion.json b/configs/scripted_insertion.json new file mode 100644 index 0000000..4e0fc94 --- /dev/null +++ b/configs/scripted_insertion.json @@ -0,0 +1,15 @@ +{ + "name": "scripted-insertion", + "inherits": "scripted-tea-bag", + "status": "fully runnable scripted-policy insertion reproduction", + "environment": { + "speed_values": [1.0, 1.5, 2.0, 2.5, 3.0] + }, + "evaluation": { + "episodes": 100, + "episodes_per_speed": 100, + "speed_start": 1.0, + "speed_stop": 3.0, + "speed_step": 0.25 + } +} diff --git a/configs/scripted_pick_and_place.json b/configs/scripted_pick_and_place.json new file mode 100644 index 0000000..a4bbd12 --- /dev/null +++ b/configs/scripted_pick_and_place.json @@ -0,0 +1,15 @@ +{ + "name": "scripted-pick-and-place", + "inherits": "scripted-tea-bag", + "status": "fully runnable scripted-policy pick-and-place reproduction", + "environment": { + "speed_values": [1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5] + }, + "evaluation": { + "episodes": 100, + "episodes_per_speed": 100, + "speed_start": 1.0, + "speed_stop": 4.5, + "speed_step": 0.25 + } +} diff --git a/configs/scripted_tea_bag.json b/configs/scripted_tea_bag.json new file mode 100644 index 0000000..25211fa --- /dev/null +++ b/configs/scripted_tea_bag.json @@ -0,0 +1,47 @@ +{ + "name": "scripted-tea-bag", + "status": "fully runnable scripted-policy reproduction", + "evidence": "Retained run_speed_rl.py at commit 097a46f plus the ICRA 2025 paper", + "base_policy": { + "base_policy": "scripted" + }, + "environment": { + "speed_values": [1.0, 1.5, 2.0, 2.5, 3.0], + "frame_skip": 10 + }, + "observation": { + "speed_observation": "state", + "include_env_state": true, + "include_qpos": true, + "include_qvel": true, + "frame_stack": 1 + }, + "reward": { + "success_bonus": 100.0, + "speed_weight": 0.01, + "speed_power": 2.0 + }, + "training": { + "decisions": 100000, + "memory_size": 1000000, + "batch_size": 128, + "learning_starts": 128, + "gradient_steps": 1, + "update_schedule": "episode", + "checkpoint_interval": 10000, + "hidden_dim": 256, + "gamma": 0.99, + "target_update": 50, + "beta_schedule": "legacy", + "epsilon": 0.0, + "min_epsilon": 0.0, + "exploration_steps": 0 + }, + "evaluation": { + "episodes": 20, + "episodes_per_speed": 50, + "speed_start": 1.0, + "speed_stop": 3.0, + "speed_step": 0.25 + } +} diff --git a/configs/scripted_tea_bag_randomized.json b/configs/scripted_tea_bag_randomized.json new file mode 100644 index 0000000..5e03f73 --- /dev/null +++ b/configs/scripted_tea_bag_randomized.json @@ -0,0 +1,12 @@ +{ + "name": "scripted-tea-bag-randomized", + "inherits": "scripted-tea-bag", + "status": "scripted-policy reproduction with seeded pose variation", + "environment": { + "randomize_object_pose": true + }, + "evaluation": { + "episodes": 100, + "episodes_per_speed": 100 + } +} diff --git a/constants.py b/constants.py new file mode 100644 index 0000000..e6f3853 --- /dev/null +++ b/constants.py @@ -0,0 +1,46 @@ +import os +from importlib import resources + +### Task parameters +### Simulation envs fixed constants +DT = 0.02 +JOINT_NAMES = ["waist", "shoulder", "elbow", "forearm_roll", "wrist_angle", "wrist_rotate"] +START_ARM_POSE = [0, -0.96, 1.16, 0, -0.3, 0, 0.02239, -0.02239, 0, -0.96, 1.16, 0, -0.3, 0, 0.02239, -0.02239] + +XML_DIR = str(resources.files("assets")) + os.sep + +# Left finger position limits (qpos[7]), right_finger = -1 * left_finger +MASTER_GRIPPER_POSITION_OPEN = 0.02417 +MASTER_GRIPPER_POSITION_CLOSE = 0.01244 +PUPPET_GRIPPER_POSITION_OPEN = 0.05800 +PUPPET_GRIPPER_POSITION_CLOSE = 0.01844 + +# Gripper joint limits (qpos[6]) +MASTER_GRIPPER_JOINT_OPEN = 0.3083 +MASTER_GRIPPER_JOINT_CLOSE = -0.6842 +PUPPET_GRIPPER_JOINT_OPEN = 1.4910 +PUPPET_GRIPPER_JOINT_CLOSE = -0.6213 + +############################ Helper functions ############################ + +MASTER_GRIPPER_POSITION_NORMALIZE_FN = lambda x: (x - MASTER_GRIPPER_POSITION_CLOSE) / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) +PUPPET_GRIPPER_POSITION_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_POSITION_CLOSE) / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) +MASTER_GRIPPER_POSITION_UNNORMALIZE_FN = lambda x: x * (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) + MASTER_GRIPPER_POSITION_CLOSE +PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN = lambda x: x * (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) + PUPPET_GRIPPER_POSITION_CLOSE +MASTER2PUPPET_POSITION_FN = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN(MASTER_GRIPPER_POSITION_NORMALIZE_FN(x)) + +MASTER_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) +PUPPET_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) +MASTER_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE +PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE +MASTER2PUPPET_JOINT_FN = lambda x: PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN(MASTER_GRIPPER_JOINT_NORMALIZE_FN(x)) + +MASTER_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) +PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) + +MASTER_POS2JOINT = lambda x: MASTER_GRIPPER_POSITION_NORMALIZE_FN(x) * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE +MASTER_JOINT2POS = lambda x: MASTER_GRIPPER_POSITION_UNNORMALIZE_FN((x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE)) +PUPPET_POS2JOINT = lambda x: PUPPET_GRIPPER_POSITION_NORMALIZE_FN(x) * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE +PUPPET_JOINT2POS = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN((x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE)) + +MASTER_GRIPPER_JOINT_MID = (MASTER_GRIPPER_JOINT_OPEN + MASTER_GRIPPER_JOINT_CLOSE)/2 diff --git a/detr/LICENSE b/detr/LICENSE new file mode 100644 index 0000000..b1395e9 --- /dev/null +++ b/detr/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2020 - present, Facebook, Inc + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/detr/README.md b/detr/README.md new file mode 100644 index 0000000..500b1b8 --- /dev/null +++ b/detr/README.md @@ -0,0 +1,9 @@ +This part of the codebase is modified from DETR https://github.com/facebookresearch/detr under APACHE 2.0. + + @article{Carion2020EndtoEndOD, + title={End-to-End Object Detection with Transformers}, + author={Nicolas Carion and Francisco Massa and Gabriel Synnaeve and Nicolas Usunier and Alexander Kirillov and Sergey Zagoruyko}, + journal={ArXiv}, + year={2020}, + volume={abs/2005.12872} + } \ No newline at end of file diff --git a/detr/__init__.py b/detr/__init__.py new file mode 100644 index 0000000..089d1af --- /dev/null +++ b/detr/__init__.py @@ -0,0 +1 @@ +"""ACT/DETR model components retained for chunk-policy compatibility.""" diff --git a/detr/main.py b/detr/main.py new file mode 100644 index 0000000..7315bc2 --- /dev/null +++ b/detr/main.py @@ -0,0 +1,77 @@ +"""Construct the retained ACT/DETR models from a dictionary configuration.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from .models import build_ACT_model, build_CNNMLP_model + + +MODEL_DEFAULTS = { + "lr": 1e-4, + "lr_backbone": 1e-5, + "weight_decay": 1e-4, + "backbone": "resnet18", + "pretrained_backbone": False, + "dilation": False, + "position_embedding": "sine", + "camera_names": [], + "enc_layers": 4, + "dec_layers": 6, + "dim_feedforward": 2048, + "hidden_dim": 256, + "dropout": 0.1, + "nheads": 8, + "num_queries": 400, + "pre_norm": False, + "masks": False, +} + + +def _model_args(overrides): + values = {**MODEL_DEFAULTS, **overrides} + return SimpleNamespace(**values) + + +def _optimizer(model, args): + parameter_groups = [ + { + "params": [ + parameter + for name, parameter in model.named_parameters() + if "backbone" not in name and parameter.requires_grad + ] + }, + { + "params": [ + parameter + for name, parameter in model.named_parameters() + if "backbone" in name and parameter.requires_grad + ], + "lr": args.lr_backbone, + }, + ] + return torch.optim.AdamW( + parameter_groups, lr=args.lr, weight_decay=args.weight_decay + ) + + +def _device(args): + return torch.device( + getattr(args, "device", None) + or ("cuda" if torch.cuda.is_available() else "cpu") + ) + + +def build_ACT_model_and_optimizer(args_override): + args = _model_args(args_override) + model = build_ACT_model(args).to(_device(args)) + return model, _optimizer(model, args) + + +def build_CNNMLP_model_and_optimizer(args_override): + args = _model_args(args_override) + model = build_CNNMLP_model(args).to(_device(args)) + return model, _optimizer(model, args) diff --git a/detr/models/__init__.py b/detr/models/__init__.py new file mode 100644 index 0000000..cc78db1 --- /dev/null +++ b/detr/models/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +from .detr_vae import build as build_vae +from .detr_vae import build_cnnmlp as build_cnnmlp + +def build_ACT_model(args): + return build_vae(args) + +def build_CNNMLP_model(args): + return build_cnnmlp(args) \ No newline at end of file diff --git a/detr/models/backbone.py b/detr/models/backbone.py new file mode 100644 index 0000000..0fc349a --- /dev/null +++ b/detr/models/backbone.py @@ -0,0 +1,127 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +Backbone modules. +""" +from collections import OrderedDict + +import torch +import torch.nn.functional as F +import torchvision +from torch import nn +from torchvision.models._utils import IntermediateLayerGetter +from typing import Dict, List + +from ..util.misc import NestedTensor, is_main_process + +from .position_encoding import build_position_encoding + +class FrozenBatchNorm2d(torch.nn.Module): + """ + BatchNorm2d where the batch statistics and the affine parameters are fixed. + + Copy-paste from torchvision.misc.ops with added eps before rqsrt, + without which any other policy_models than torchvision.policy_models.resnet[18,34,50,101] + produce nans. + """ + + def __init__(self, n): + super(FrozenBatchNorm2d, self).__init__() + self.register_buffer("weight", torch.ones(n)) + self.register_buffer("bias", torch.zeros(n)) + self.register_buffer("running_mean", torch.zeros(n)) + self.register_buffer("running_var", torch.ones(n)) + + def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, + missing_keys, unexpected_keys, error_msgs): + num_batches_tracked_key = prefix + 'num_batches_tracked' + if num_batches_tracked_key in state_dict: + del state_dict[num_batches_tracked_key] + + super(FrozenBatchNorm2d, self)._load_from_state_dict( + state_dict, prefix, local_metadata, strict, + missing_keys, unexpected_keys, error_msgs) + + def forward(self, x): + # move reshapes to the beginning + # to make it fuser-friendly + w = self.weight.reshape(1, -1, 1, 1) + b = self.bias.reshape(1, -1, 1, 1) + rv = self.running_var.reshape(1, -1, 1, 1) + rm = self.running_mean.reshape(1, -1, 1, 1) + eps = 1e-5 + scale = w * (rv + eps).rsqrt() + bias = b - rm * scale + return x * scale + bias + + +class BackboneBase(nn.Module): + + def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool): + super().__init__() + # for name, parameter in backbone.named_parameters(): # only train later layers # TODO do we want this? + # if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name: + # parameter.requires_grad_(False) + if return_interm_layers: + return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"} + else: + return_layers = {'layer4': "0"} + self.body = IntermediateLayerGetter(backbone, return_layers=return_layers) + self.num_channels = num_channels + + def forward(self, tensor): + xs = self.body(tensor) + return xs + # out: Dict[str, NestedTensor] = {} + # for name, x in xs.items(): + # m = tensor_list.mask + # assert m is not None + # mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0] + # out[name] = NestedTensor(x, mask) + # return out + + +class Backbone(BackboneBase): + """ResNet backbone with frozen BatchNorm.""" + def __init__(self, name: str, + train_backbone: bool, + return_interm_layers: bool, + dilation: bool, + pretrained: bool = True): + backbone = getattr(torchvision.models, name)( + replace_stride_with_dilation=[False, False, dilation], + weights="DEFAULT" if pretrained and is_main_process() else None, + norm_layer=FrozenBatchNorm2d) + num_channels = 512 if name in ('resnet18', 'resnet34') else 2048 + super().__init__(backbone, train_backbone, num_channels, return_interm_layers) + + +class Joiner(nn.Sequential): + def __init__(self, backbone, position_embedding): + super().__init__(backbone, position_embedding) + + def forward(self, tensor_list: NestedTensor): + xs = self[0](tensor_list) + out: List[NestedTensor] = [] + pos = [] + for name, x in xs.items(): + out.append(x) + # position encoding + pos.append(self[1](x).to(x.dtype)) + + return out, pos + + +def build_backbone(args): + position_embedding = build_position_encoding(args) + train_backbone = args.lr_backbone > 0 + return_interm_layers = args.masks + backbone = Backbone( + args.backbone, + train_backbone, + return_interm_layers, + args.dilation, + pretrained=getattr(args, "pretrained_backbone", True), + ) + model = Joiner(backbone, position_embedding) + model.num_channels = backbone.num_channels + return model diff --git a/detr/models/detr_vae.py b/detr/models/detr_vae.py new file mode 100644 index 0000000..69e999c --- /dev/null +++ b/detr/models/detr_vae.py @@ -0,0 +1,274 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +DETR model and criterion classes. +""" +import torch +from torch import nn +from torch.autograd import Variable +from .backbone import build_backbone +from .transformer import build_transformer, TransformerEncoder, TransformerEncoderLayer + +import numpy as np + + +def reparametrize(mu, logvar): + std = logvar.div(2).exp() + eps = Variable(std.data.new(std.size()).normal_()) + return mu + std * eps + + +def get_sinusoid_encoding_table(n_position, d_hid): + def get_position_angle_vec(position): + return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)] + + sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)]) + sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i + sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1 + + return torch.FloatTensor(sinusoid_table).unsqueeze(0) + + +class DETRVAE(nn.Module): + """ This is the DETR module that performs object detection """ + def __init__(self, backbones, transformer, encoder, state_dim, num_queries, camera_names): + """ Initializes the model. + Parameters: + backbones: torch module of the backbone to be used. See backbone.py + transformer: torch module of the transformer architecture. See transformer.py + state_dim: robot state dimension of the environment + num_queries: number of object queries, ie detection slot. This is the maximal number of objects + DETR can detect in a single image. For COCO, we recommend 100 queries. + aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used. + """ + super().__init__() + self.num_queries = num_queries + self.camera_names = camera_names + self.transformer = transformer + self.encoder = encoder + hidden_dim = transformer.d_model + self.action_head = nn.Linear(hidden_dim, state_dim) + self.is_pad_head = nn.Linear(hidden_dim, 1) + self.query_embed = nn.Embedding(num_queries, hidden_dim) + if backbones is not None: + self.input_proj = nn.Conv2d(backbones[0].num_channels, hidden_dim, kernel_size=1) + self.backbones = nn.ModuleList(backbones) + self.input_proj_robot_state = nn.Linear(14, hidden_dim) + else: + # input_dim = 14 + 7 # robot_state + env_state + self.input_proj_robot_state = nn.Linear(14, hidden_dim) + self.input_proj_env_state = nn.Linear(7, hidden_dim) + self.pos = torch.nn.Embedding(2, hidden_dim) + self.backbones = None + + # encoder extra parameters + self.latent_dim = 32 # final size of latent z # TODO tune + self.cls_embed = nn.Embedding(1, hidden_dim) # extra cls token embedding + self.encoder_action_proj = nn.Linear(14, hidden_dim) # project action to embedding + self.encoder_joint_proj = nn.Linear(14, hidden_dim) # project qpos to embedding + self.latent_proj = nn.Linear(hidden_dim, self.latent_dim*2) # project hidden state to latent std, var + self.register_buffer('pos_table', get_sinusoid_encoding_table(1+1+num_queries, hidden_dim)) # [CLS], qpos, a_seq + + # decoder extra parameters + self.latent_out_proj = nn.Linear(self.latent_dim, hidden_dim) # project latent sample to embedding + self.additional_pos_embed = nn.Embedding(2, hidden_dim) # learned position embedding for proprio and latent + + def forward(self, qpos, image, env_state, actions=None, is_pad=None): + """ + qpos: batch, qpos_dim + image: batch, num_cam, channel, height, width + env_state: None + actions: batch, seq, action_dim + """ + is_training = actions is not None # train or val + bs, _ = qpos.shape + ### Obtain latent z from action sequence + if is_training: + # project action sequence to embedding dim, and concat with a CLS token + action_embed = self.encoder_action_proj(actions) # (bs, seq, hidden_dim) + qpos_embed = self.encoder_joint_proj(qpos) # (bs, hidden_dim) + qpos_embed = torch.unsqueeze(qpos_embed, axis=1) # (bs, 1, hidden_dim) + cls_embed = self.cls_embed.weight # (1, hidden_dim) + cls_embed = torch.unsqueeze(cls_embed, axis=0).repeat(bs, 1, 1) # (bs, 1, hidden_dim) + encoder_input = torch.cat([cls_embed, qpos_embed, action_embed], axis=1) # (bs, seq+1, hidden_dim) + encoder_input = encoder_input.permute(1, 0, 2) # (seq+1, bs, hidden_dim) + # do not mask cls token + cls_joint_is_pad = torch.full((bs, 2), False).to(qpos.device) # False: not a padding + is_pad = torch.cat([cls_joint_is_pad, is_pad], axis=1) # (bs, seq+1) + # obtain position embedding + pos_embed = self.pos_table.clone().detach() + pos_embed = pos_embed.permute(1, 0, 2) # (seq+1, 1, hidden_dim) + # query model + encoder_output = self.encoder(encoder_input, pos=pos_embed, src_key_padding_mask=is_pad) + encoder_output = encoder_output[0] # take cls output only + latent_info = self.latent_proj(encoder_output) + mu = latent_info[:, :self.latent_dim] + logvar = latent_info[:, self.latent_dim:] + latent_sample = reparametrize(mu, logvar) + latent_input = self.latent_out_proj(latent_sample) + else: + mu = logvar = None + latent_sample = torch.zeros([bs, self.latent_dim], dtype=torch.float32).to(qpos.device) + latent_input = self.latent_out_proj(latent_sample) + + if self.backbones is not None: + # Image observation features and position embeddings + all_cam_features = [] + all_cam_pos = [] + for cam_id, cam_name in enumerate(self.camera_names): + features, pos = self.backbones[0](image[:, cam_id]) # HARDCODED + features = features[0] # take the last layer feature + pos = pos[0] + all_cam_features.append(self.input_proj(features)) + all_cam_pos.append(pos) + # proprioception features + proprio_input = self.input_proj_robot_state(qpos) + # fold camera dimension into width dimension + src = torch.cat(all_cam_features, axis=3) + pos = torch.cat(all_cam_pos, axis=3) + hs = self.transformer(src, None, self.query_embed.weight, pos, latent_input, proprio_input, self.additional_pos_embed.weight)[0] + else: + qpos = self.input_proj_robot_state(qpos) + env_state = self.input_proj_env_state(env_state) + transformer_input = torch.cat([qpos, env_state], axis=1) # seq length = 2 + hs = self.transformer(transformer_input, None, self.query_embed.weight, self.pos.weight)[0] + a_hat = self.action_head(hs) + is_pad_hat = self.is_pad_head(hs) + return a_hat, is_pad_hat, [mu, logvar] + + + +class CNNMLP(nn.Module): + def __init__(self, backbones, state_dim, camera_names): + """ Initializes the model. + Parameters: + backbones: torch module of the backbone to be used. See backbone.py + transformer: torch module of the transformer architecture. See transformer.py + state_dim: robot state dimension of the environment + num_queries: number of object queries, ie detection slot. This is the maximal number of objects + DETR can detect in a single image. For COCO, we recommend 100 queries. + aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used. + """ + super().__init__() + self.camera_names = camera_names + self.action_head = nn.Linear(1000, state_dim) # TODO add more + if backbones is not None: + self.backbones = nn.ModuleList(backbones) + backbone_down_projs = [] + for backbone in backbones: + down_proj = nn.Sequential( + nn.Conv2d(backbone.num_channels, 128, kernel_size=5), + nn.Conv2d(128, 64, kernel_size=5), + nn.Conv2d(64, 32, kernel_size=5) + ) + backbone_down_projs.append(down_proj) + self.backbone_down_projs = nn.ModuleList(backbone_down_projs) + + mlp_in_dim = 768 * len(backbones) + 14 + self.mlp = mlp(input_dim=mlp_in_dim, hidden_dim=1024, output_dim=14, hidden_depth=2) + else: + raise NotImplementedError + + def forward(self, qpos, image, env_state, actions=None): + """ + qpos: batch, qpos_dim + image: batch, num_cam, channel, height, width + env_state: None + actions: batch, seq, action_dim + """ + is_training = actions is not None # train or val + bs, _ = qpos.shape + # Image observation features and position embeddings + all_cam_features = [] + for cam_id, cam_name in enumerate(self.camera_names): + features, pos = self.backbones[cam_id](image[:, cam_id]) + features = features[0] # take the last layer feature + pos = pos[0] # not used + all_cam_features.append(self.backbone_down_projs[cam_id](features)) + # flatten everything + flattened_features = [] + for cam_feature in all_cam_features: + flattened_features.append(cam_feature.reshape([bs, -1])) + flattened_features = torch.cat(flattened_features, axis=1) # 768 each + features = torch.cat([flattened_features, qpos], axis=1) # qpos: 14 + a_hat = self.mlp(features) + return a_hat + + +def mlp(input_dim, hidden_dim, output_dim, hidden_depth): + if hidden_depth == 0: + mods = [nn.Linear(input_dim, output_dim)] + else: + mods = [nn.Linear(input_dim, hidden_dim), nn.ReLU(inplace=True)] + for i in range(hidden_depth - 1): + mods += [nn.Linear(hidden_dim, hidden_dim), nn.ReLU(inplace=True)] + mods.append(nn.Linear(hidden_dim, output_dim)) + trunk = nn.Sequential(*mods) + return trunk + + +def build_encoder(args): + d_model = args.hidden_dim # 256 + dropout = args.dropout # 0.1 + nhead = args.nheads # 8 + dim_feedforward = args.dim_feedforward # 2048 + num_encoder_layers = args.enc_layers # 4 # TODO shared with VAE decoder + normalize_before = args.pre_norm # False + activation = "relu" + + encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, + dropout, activation, normalize_before) + encoder_norm = nn.LayerNorm(d_model) if normalize_before else None + encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm) + + return encoder + + +def build(args): + state_dim = 14 # TODO hardcode + + # From state + # backbone = None # from state for now, no need for conv nets + # From image + backbones = [] + backbone = build_backbone(args) + backbones.append(backbone) + + transformer = build_transformer(args) + + encoder = build_encoder(args) + + model = DETRVAE( + backbones, + transformer, + encoder, + state_dim=state_dim, + num_queries=args.num_queries, + camera_names=args.camera_names, + ) + + n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) + print("number of parameters: %.2fM" % (n_parameters/1e6,)) + + return model + +def build_cnnmlp(args): + state_dim = 14 # TODO hardcode + + # From state + # backbone = None # from state for now, no need for conv nets + # From image + backbones = [] + for _ in args.camera_names: + backbone = build_backbone(args) + backbones.append(backbone) + + model = CNNMLP( + backbones, + state_dim=state_dim, + camera_names=args.camera_names, + ) + + n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad) + print("number of parameters: %.2fM" % (n_parameters/1e6,)) + + return model diff --git a/detr/models/position_encoding.py b/detr/models/position_encoding.py new file mode 100644 index 0000000..db455d9 --- /dev/null +++ b/detr/models/position_encoding.py @@ -0,0 +1,90 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +Various positional encodings for the transformer. +""" +import math +import torch +from torch import nn + +from ..util.misc import NestedTensor + +class PositionEmbeddingSine(nn.Module): + """ + This is a more standard version of the position embedding, very similar to the one + used by the Attention is all you need paper, generalized to work on images. + """ + def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None): + super().__init__() + self.num_pos_feats = num_pos_feats + self.temperature = temperature + self.normalize = normalize + if scale is not None and normalize is False: + raise ValueError("normalize should be True if scale is passed") + if scale is None: + scale = 2 * math.pi + self.scale = scale + + def forward(self, tensor): + x = tensor + # mask = tensor_list.mask + # assert mask is not None + # not_mask = ~mask + + not_mask = torch.ones_like(x[0, [0]]) + y_embed = not_mask.cumsum(1, dtype=torch.float32) + x_embed = not_mask.cumsum(2, dtype=torch.float32) + if self.normalize: + eps = 1e-6 + y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale + x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale + + dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device) + dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats) + + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) + pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) + pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + return pos + + +class PositionEmbeddingLearned(nn.Module): + """ + Absolute pos embedding, learned. + """ + def __init__(self, num_pos_feats=256): + super().__init__() + self.row_embed = nn.Embedding(50, num_pos_feats) + self.col_embed = nn.Embedding(50, num_pos_feats) + self.reset_parameters() + + def reset_parameters(self): + nn.init.uniform_(self.row_embed.weight) + nn.init.uniform_(self.col_embed.weight) + + def forward(self, tensor_list: NestedTensor): + x = tensor_list.tensors + h, w = x.shape[-2:] + i = torch.arange(w, device=x.device) + j = torch.arange(h, device=x.device) + x_emb = self.col_embed(i) + y_emb = self.row_embed(j) + pos = torch.cat([ + x_emb.unsqueeze(0).repeat(h, 1, 1), + y_emb.unsqueeze(1).repeat(1, w, 1), + ], dim=-1).permute(2, 0, 1).unsqueeze(0).repeat(x.shape[0], 1, 1, 1) + return pos + + +def build_position_encoding(args): + N_steps = args.hidden_dim // 2 + if args.position_embedding in ('v2', 'sine'): + # TODO find a better way of exposing other arguments + position_embedding = PositionEmbeddingSine(N_steps, normalize=True) + elif args.position_embedding in ('v3', 'learned'): + position_embedding = PositionEmbeddingLearned(N_steps) + else: + raise ValueError(f"not supported {args.position_embedding}") + + return position_embedding diff --git a/detr/models/transformer.py b/detr/models/transformer.py new file mode 100644 index 0000000..f14278e --- /dev/null +++ b/detr/models/transformer.py @@ -0,0 +1,311 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +DETR Transformer class. + +Copy-paste from torch.nn.Transformer with modifications: + * positional encodings are passed in MHattention + * extra LN at the end of encoder is removed + * decoder returns a stack of activations from all decoding layers +""" +import copy +from typing import Optional, List + +import torch +import torch.nn.functional as F +from torch import nn, Tensor + +class Transformer(nn.Module): + + def __init__(self, d_model=512, nhead=8, num_encoder_layers=6, + num_decoder_layers=6, dim_feedforward=2048, dropout=0.1, + activation="relu", normalize_before=False, + return_intermediate_dec=False): + super().__init__() + + encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, + dropout, activation, normalize_before) + encoder_norm = nn.LayerNorm(d_model) if normalize_before else None + self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm) + + decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, + dropout, activation, normalize_before) + decoder_norm = nn.LayerNorm(d_model) + self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm, + return_intermediate=return_intermediate_dec) + + self._reset_parameters() + + self.d_model = d_model + self.nhead = nhead + + def _reset_parameters(self): + for p in self.parameters(): + if p.dim() > 1: + nn.init.xavier_uniform_(p) + + def forward(self, src, mask, query_embed, pos_embed, latent_input=None, proprio_input=None, additional_pos_embed=None): + # TODO flatten only when input has H and W + if len(src.shape) == 4: # has H and W + # flatten NxCxHxW to HWxNxC + bs, c, h, w = src.shape + src = src.flatten(2).permute(2, 0, 1) + pos_embed = pos_embed.flatten(2).permute(2, 0, 1).repeat(1, bs, 1) + query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1) + # mask = mask.flatten(1) + + additional_pos_embed = additional_pos_embed.unsqueeze(1).repeat(1, bs, 1) # seq, bs, dim + pos_embed = torch.cat([additional_pos_embed, pos_embed], axis=0) + + addition_input = torch.stack([latent_input, proprio_input], axis=0) + src = torch.cat([addition_input, src], axis=0) + else: + assert len(src.shape) == 3 + # flatten NxHWxC to HWxNxC + bs, hw, c = src.shape + src = src.permute(1, 0, 2) + pos_embed = pos_embed.unsqueeze(1).repeat(1, bs, 1) + query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1) + + tgt = torch.zeros_like(query_embed) + memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed) + hs = self.decoder(tgt, memory, memory_key_padding_mask=mask, + pos=pos_embed, query_pos=query_embed) + hs = hs.transpose(1, 2) + return hs + +class TransformerEncoder(nn.Module): + + def __init__(self, encoder_layer, num_layers, norm=None): + super().__init__() + self.layers = _get_clones(encoder_layer, num_layers) + self.num_layers = num_layers + self.norm = norm + + def forward(self, src, + mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None): + output = src + + for layer in self.layers: + output = layer(output, src_mask=mask, + src_key_padding_mask=src_key_padding_mask, pos=pos) + + if self.norm is not None: + output = self.norm(output) + + return output + + +class TransformerDecoder(nn.Module): + + def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False): + super().__init__() + self.layers = _get_clones(decoder_layer, num_layers) + self.num_layers = num_layers + self.norm = norm + self.return_intermediate = return_intermediate + + def forward(self, tgt, memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + output = tgt + + intermediate = [] + + for layer in self.layers: + output = layer(output, memory, tgt_mask=tgt_mask, + memory_mask=memory_mask, + tgt_key_padding_mask=tgt_key_padding_mask, + memory_key_padding_mask=memory_key_padding_mask, + pos=pos, query_pos=query_pos) + if self.return_intermediate: + intermediate.append(self.norm(output)) + + if self.norm is not None: + output = self.norm(output) + if self.return_intermediate: + intermediate.pop() + intermediate.append(output) + + if self.return_intermediate: + return torch.stack(intermediate) + + return output.unsqueeze(0) + + +class TransformerEncoderLayer(nn.Module): + + def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, + activation="relu", normalize_before=False): + super().__init__() + self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + self.normalize_before = normalize_before + + def with_pos_embed(self, tensor, pos: Optional[Tensor]): + return tensor if pos is None else tensor + pos + + def forward_post(self, + src, + src_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None): + q = k = self.with_pos_embed(src, pos) + src2 = self.self_attn(q, k, value=src, attn_mask=src_mask, + key_padding_mask=src_key_padding_mask)[0] + src = src + self.dropout1(src2) + src = self.norm1(src) + src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) + src = src + self.dropout2(src2) + src = self.norm2(src) + return src + + def forward_pre(self, src, + src_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None): + src2 = self.norm1(src) + q = k = self.with_pos_embed(src2, pos) + src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask, + key_padding_mask=src_key_padding_mask)[0] + src = src + self.dropout1(src2) + src2 = self.norm2(src) + src2 = self.linear2(self.dropout(self.activation(self.linear1(src2)))) + src = src + self.dropout2(src2) + return src + + def forward(self, src, + src_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None): + if self.normalize_before: + return self.forward_pre(src, src_mask, src_key_padding_mask, pos) + return self.forward_post(src, src_mask, src_key_padding_mask, pos) + + +class TransformerDecoderLayer(nn.Module): + + def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, + activation="relu", normalize_before=False): + super().__init__() + self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) + # Implementation of Feedforward model + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + + self.norm1 = nn.LayerNorm(d_model) + self.norm2 = nn.LayerNorm(d_model) + self.norm3 = nn.LayerNorm(d_model) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + self.dropout3 = nn.Dropout(dropout) + + self.activation = _get_activation_fn(activation) + self.normalize_before = normalize_before + + def with_pos_embed(self, tensor, pos: Optional[Tensor]): + return tensor if pos is None else tensor + pos + + def forward_post(self, tgt, memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + q = k = self.with_pos_embed(tgt, query_pos) + tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask, + key_padding_mask=tgt_key_padding_mask)[0] + tgt = tgt + self.dropout1(tgt2) + tgt = self.norm1(tgt) + tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos), + key=self.with_pos_embed(memory, pos), + value=memory, attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask)[0] + tgt = tgt + self.dropout2(tgt2) + tgt = self.norm2(tgt) + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) + tgt = tgt + self.dropout3(tgt2) + tgt = self.norm3(tgt) + return tgt + + def forward_pre(self, tgt, memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + tgt2 = self.norm1(tgt) + q = k = self.with_pos_embed(tgt2, query_pos) + tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask, + key_padding_mask=tgt_key_padding_mask)[0] + tgt = tgt + self.dropout1(tgt2) + tgt2 = self.norm2(tgt) + tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos), + key=self.with_pos_embed(memory, pos), + value=memory, attn_mask=memory_mask, + key_padding_mask=memory_key_padding_mask)[0] + tgt = tgt + self.dropout2(tgt2) + tgt2 = self.norm3(tgt) + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) + tgt = tgt + self.dropout3(tgt2) + return tgt + + def forward(self, tgt, memory, + tgt_mask: Optional[Tensor] = None, + memory_mask: Optional[Tensor] = None, + tgt_key_padding_mask: Optional[Tensor] = None, + memory_key_padding_mask: Optional[Tensor] = None, + pos: Optional[Tensor] = None, + query_pos: Optional[Tensor] = None): + if self.normalize_before: + return self.forward_pre(tgt, memory, tgt_mask, memory_mask, + tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos) + return self.forward_post(tgt, memory, tgt_mask, memory_mask, + tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos) + + +def _get_clones(module, N): + return nn.ModuleList([copy.deepcopy(module) for i in range(N)]) + + +def build_transformer(args): + return Transformer( + d_model=args.hidden_dim, + dropout=args.dropout, + nhead=args.nheads, + dim_feedforward=args.dim_feedforward, + num_encoder_layers=args.enc_layers, + num_decoder_layers=args.dec_layers, + normalize_before=args.pre_norm, + return_intermediate_dec=True, + ) + + +def _get_activation_fn(activation): + """Return an activation function given a string""" + if activation == "relu": + return F.relu + if activation == "gelu": + return F.gelu + if activation == "glu": + return F.glu + raise RuntimeError(F"activation should be relu/gelu, not {activation}.") diff --git a/detr/util/__init__.py b/detr/util/__init__.py new file mode 100644 index 0000000..168f997 --- /dev/null +++ b/detr/util/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved diff --git a/detr/util/misc.py b/detr/util/misc.py new file mode 100644 index 0000000..dfa9fb5 --- /dev/null +++ b/detr/util/misc.py @@ -0,0 +1,468 @@ +# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved +""" +Misc functions, including distributed helpers. + +Mostly copy-paste from torchvision references. +""" +import os +import subprocess +import time +from collections import defaultdict, deque +import datetime +import pickle +from packaging import version +from typing import Optional, List + +import torch +import torch.distributed as dist +from torch import Tensor + +# needed due to empty tensor bug in pytorch and torchvision 0.5 +import torchvision +if version.parse(torchvision.__version__) < version.parse('0.7'): + from torchvision.ops import _new_empty_tensor + from torchvision.ops.misc import _output_size + + +class SmoothedValue(object): + """Track a series of values and provide access to smoothed values over a + window or the global series average. + """ + + def __init__(self, window_size=20, fmt=None): + if fmt is None: + fmt = "{median:.4f} ({global_avg:.4f})" + self.deque = deque(maxlen=window_size) + self.total = 0.0 + self.count = 0 + self.fmt = fmt + + def update(self, value, n=1): + self.deque.append(value) + self.count += n + self.total += value * n + + def synchronize_between_processes(self): + """ + Warning: does not synchronize the deque! + """ + if not is_dist_avail_and_initialized(): + return + t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda') + dist.barrier() + dist.all_reduce(t) + t = t.tolist() + self.count = int(t[0]) + self.total = t[1] + + @property + def median(self): + d = torch.tensor(list(self.deque)) + return d.median().item() + + @property + def avg(self): + d = torch.tensor(list(self.deque), dtype=torch.float32) + return d.mean().item() + + @property + def global_avg(self): + return self.total / self.count + + @property + def max(self): + return max(self.deque) + + @property + def value(self): + return self.deque[-1] + + def __str__(self): + return self.fmt.format( + median=self.median, + avg=self.avg, + global_avg=self.global_avg, + max=self.max, + value=self.value) + + +def all_gather(data): + """ + Run all_gather on arbitrary picklable data (not necessarily tensors) + Args: + data: any picklable object + Returns: + list[data]: list of data gathered from each rank + """ + world_size = get_world_size() + if world_size == 1: + return [data] + + # serialized to a Tensor + buffer = pickle.dumps(data) + storage = torch.ByteStorage.from_buffer(buffer) + tensor = torch.ByteTensor(storage).to("cuda") + + # obtain Tensor size of each rank + local_size = torch.tensor([tensor.numel()], device="cuda") + size_list = [torch.tensor([0], device="cuda") for _ in range(world_size)] + dist.all_gather(size_list, local_size) + size_list = [int(size.item()) for size in size_list] + max_size = max(size_list) + + # receiving Tensor from all ranks + # we pad the tensor because torch all_gather does not support + # gathering tensors of different shapes + tensor_list = [] + for _ in size_list: + tensor_list.append(torch.empty((max_size,), dtype=torch.uint8, device="cuda")) + if local_size != max_size: + padding = torch.empty(size=(max_size - local_size,), dtype=torch.uint8, device="cuda") + tensor = torch.cat((tensor, padding), dim=0) + dist.all_gather(tensor_list, tensor) + + data_list = [] + for size, tensor in zip(size_list, tensor_list): + buffer = tensor.cpu().numpy().tobytes()[:size] + data_list.append(pickle.loads(buffer)) + + return data_list + + +def reduce_dict(input_dict, average=True): + """ + Args: + input_dict (dict): all the values will be reduced + average (bool): whether to do average or sum + Reduce the values in the dictionary from all processes so that all processes + have the averaged results. Returns a dict with the same fields as + input_dict, after reduction. + """ + world_size = get_world_size() + if world_size < 2: + return input_dict + with torch.no_grad(): + names = [] + values = [] + # sort the keys so that they are consistent across processes + for k in sorted(input_dict.keys()): + names.append(k) + values.append(input_dict[k]) + values = torch.stack(values, dim=0) + dist.all_reduce(values) + if average: + values /= world_size + reduced_dict = {k: v for k, v in zip(names, values)} + return reduced_dict + + +class MetricLogger(object): + def __init__(self, delimiter="\t"): + self.meters = defaultdict(SmoothedValue) + self.delimiter = delimiter + + def update(self, **kwargs): + for k, v in kwargs.items(): + if isinstance(v, torch.Tensor): + v = v.item() + assert isinstance(v, (float, int)) + self.meters[k].update(v) + + def __getattr__(self, attr): + if attr in self.meters: + return self.meters[attr] + if attr in self.__dict__: + return self.__dict__[attr] + raise AttributeError("'{}' object has no attribute '{}'".format( + type(self).__name__, attr)) + + def __str__(self): + loss_str = [] + for name, meter in self.meters.items(): + loss_str.append( + "{}: {}".format(name, str(meter)) + ) + return self.delimiter.join(loss_str) + + def synchronize_between_processes(self): + for meter in self.meters.values(): + meter.synchronize_between_processes() + + def add_meter(self, name, meter): + self.meters[name] = meter + + def log_every(self, iterable, print_freq, header=None): + i = 0 + if not header: + header = '' + start_time = time.time() + end = time.time() + iter_time = SmoothedValue(fmt='{avg:.4f}') + data_time = SmoothedValue(fmt='{avg:.4f}') + space_fmt = ':' + str(len(str(len(iterable)))) + 'd' + if torch.cuda.is_available(): + log_msg = self.delimiter.join([ + header, + '[{0' + space_fmt + '}/{1}]', + 'eta: {eta}', + '{meters}', + 'time: {time}', + 'data: {data}', + 'max mem: {memory:.0f}' + ]) + else: + log_msg = self.delimiter.join([ + header, + '[{0' + space_fmt + '}/{1}]', + 'eta: {eta}', + '{meters}', + 'time: {time}', + 'data: {data}' + ]) + MB = 1024.0 * 1024.0 + for obj in iterable: + data_time.update(time.time() - end) + yield obj + iter_time.update(time.time() - end) + if i % print_freq == 0 or i == len(iterable) - 1: + eta_seconds = iter_time.global_avg * (len(iterable) - i) + eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) + if torch.cuda.is_available(): + print(log_msg.format( + i, len(iterable), eta=eta_string, + meters=str(self), + time=str(iter_time), data=str(data_time), + memory=torch.cuda.max_memory_allocated() / MB)) + else: + print(log_msg.format( + i, len(iterable), eta=eta_string, + meters=str(self), + time=str(iter_time), data=str(data_time))) + i += 1 + end = time.time() + total_time = time.time() - start_time + total_time_str = str(datetime.timedelta(seconds=int(total_time))) + print('{} Total time: {} ({:.4f} s / it)'.format( + header, total_time_str, total_time / len(iterable))) + + +def get_sha(): + cwd = os.path.dirname(os.path.abspath(__file__)) + + def _run(command): + return subprocess.check_output(command, cwd=cwd).decode('ascii').strip() + sha = 'N/A' + diff = "clean" + branch = 'N/A' + try: + sha = _run(['git', 'rev-parse', 'HEAD']) + subprocess.check_output(['git', 'diff'], cwd=cwd) + diff = _run(['git', 'diff-index', 'HEAD']) + diff = "has uncommited changes" if diff else "clean" + branch = _run(['git', 'rev-parse', '--abbrev-ref', 'HEAD']) + except Exception: + pass + message = f"sha: {sha}, status: {diff}, branch: {branch}" + return message + + +def collate_fn(batch): + batch = list(zip(*batch)) + batch[0] = nested_tensor_from_tensor_list(batch[0]) + return tuple(batch) + + +def _max_by_axis(the_list): + # type: (List[List[int]]) -> List[int] + maxes = the_list[0] + for sublist in the_list[1:]: + for index, item in enumerate(sublist): + maxes[index] = max(maxes[index], item) + return maxes + + +class NestedTensor(object): + def __init__(self, tensors, mask: Optional[Tensor]): + self.tensors = tensors + self.mask = mask + + def to(self, device): + # type: (Device) -> NestedTensor # noqa + cast_tensor = self.tensors.to(device) + mask = self.mask + if mask is not None: + assert mask is not None + cast_mask = mask.to(device) + else: + cast_mask = None + return NestedTensor(cast_tensor, cast_mask) + + def decompose(self): + return self.tensors, self.mask + + def __repr__(self): + return str(self.tensors) + + +def nested_tensor_from_tensor_list(tensor_list: List[Tensor]): + # TODO make this more general + if tensor_list[0].ndim == 3: + if torchvision._is_tracing(): + # nested_tensor_from_tensor_list() does not export well to ONNX + # call _onnx_nested_tensor_from_tensor_list() instead + return _onnx_nested_tensor_from_tensor_list(tensor_list) + + # TODO make it support different-sized images + max_size = _max_by_axis([list(img.shape) for img in tensor_list]) + # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list])) + batch_shape = [len(tensor_list)] + max_size + b, c, h, w = batch_shape + dtype = tensor_list[0].dtype + device = tensor_list[0].device + tensor = torch.zeros(batch_shape, dtype=dtype, device=device) + mask = torch.ones((b, h, w), dtype=torch.bool, device=device) + for img, pad_img, m in zip(tensor_list, tensor, mask): + pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) + m[: img.shape[1], :img.shape[2]] = False + else: + raise ValueError('not supported') + return NestedTensor(tensor, mask) + + +# _onnx_nested_tensor_from_tensor_list() is an implementation of +# nested_tensor_from_tensor_list() that is supported by ONNX tracing. +@torch.jit.unused +def _onnx_nested_tensor_from_tensor_list(tensor_list: List[Tensor]) -> NestedTensor: + max_size = [] + for i in range(tensor_list[0].dim()): + max_size_i = torch.max(torch.stack([img.shape[i] for img in tensor_list]).to(torch.float32)).to(torch.int64) + max_size.append(max_size_i) + max_size = tuple(max_size) + + # work around for + # pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img) + # m[: img.shape[1], :img.shape[2]] = False + # which is not yet supported in onnx + padded_imgs = [] + padded_masks = [] + for img in tensor_list: + padding = [(s1 - s2) for s1, s2 in zip(max_size, tuple(img.shape))] + padded_img = torch.nn.functional.pad(img, (0, padding[2], 0, padding[1], 0, padding[0])) + padded_imgs.append(padded_img) + + m = torch.zeros_like(img[0], dtype=torch.int, device=img.device) + padded_mask = torch.nn.functional.pad(m, (0, padding[2], 0, padding[1]), "constant", 1) + padded_masks.append(padded_mask.to(torch.bool)) + + tensor = torch.stack(padded_imgs) + mask = torch.stack(padded_masks) + + return NestedTensor(tensor, mask=mask) + + +def setup_for_distributed(is_master): + """ + This function disables printing when not in master process + """ + import builtins as __builtin__ + builtin_print = __builtin__.print + + def print(*args, **kwargs): + force = kwargs.pop('force', False) + if is_master or force: + builtin_print(*args, **kwargs) + + __builtin__.print = print + + +def is_dist_avail_and_initialized(): + if not dist.is_available(): + return False + if not dist.is_initialized(): + return False + return True + + +def get_world_size(): + if not is_dist_avail_and_initialized(): + return 1 + return dist.get_world_size() + + +def get_rank(): + if not is_dist_avail_and_initialized(): + return 0 + return dist.get_rank() + + +def is_main_process(): + return get_rank() == 0 + + +def save_on_master(*args, **kwargs): + if is_main_process(): + torch.save(*args, **kwargs) + + +def init_distributed_mode(args): + if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ: + args.rank = int(os.environ["RANK"]) + args.world_size = int(os.environ['WORLD_SIZE']) + args.gpu = int(os.environ['LOCAL_RANK']) + elif 'SLURM_PROCID' in os.environ: + args.rank = int(os.environ['SLURM_PROCID']) + args.gpu = args.rank % torch.cuda.device_count() + else: + print('Not using distributed mode') + args.distributed = False + return + + args.distributed = True + + torch.cuda.set_device(args.gpu) + args.dist_backend = 'nccl' + print('| distributed init (rank {}): {}'.format( + args.rank, args.dist_url), flush=True) + torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url, + world_size=args.world_size, rank=args.rank) + torch.distributed.barrier() + setup_for_distributed(args.rank == 0) + + +@torch.no_grad() +def accuracy(output, target, topk=(1,)): + """Computes the precision@k for the specified values of k""" + if target.numel() == 0: + return [torch.zeros([], device=output.device)] + maxk = max(topk) + batch_size = target.size(0) + + _, pred = output.topk(maxk, 1, True, True) + pred = pred.t() + correct = pred.eq(target.view(1, -1).expand_as(pred)) + + res = [] + for k in topk: + correct_k = correct[:k].view(-1).float().sum(0) + res.append(correct_k.mul_(100.0 / batch_size)) + return res + + +def interpolate(input, size=None, scale_factor=None, mode="nearest", align_corners=None): + # type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor + """ + Equivalent to nn.functional.interpolate, but with support for empty batch sizes. + This will eventually be supported natively by PyTorch, and this + class can go away. + """ + if version.parse(torchvision.__version__) < version.parse('0.7'): + if input.numel() > 0: + return torch.nn.functional.interpolate( + input, size, scale_factor, mode, align_corners + ) + + output_shape = _output_size(2, input, size, scale_factor) + output_shape = list(input.shape[:-2]) + list(output_shape) + return _new_empty_tensor(input, output_shape) + else: + return torchvision.ops.misc.interpolate(input, size, scale_factor, mode, align_corners) diff --git a/docs/EXTERNAL_POLICIES.md b/docs/EXTERNAL_POLICIES.md new file mode 100644 index 0000000..d85f0f0 --- /dev/null +++ b/docs/EXTERNAL_POLICIES.md @@ -0,0 +1,147 @@ +# External task-policy integration + +SpeedTuning treats the task policy and speed policy as independent components. +An external task policy supplies robot action chunks; the speed policy chooses +how quickly those chunks are executed. + +This integration is optional. The bundled scripted policies are sufficient for +the from-scratch simulation reproduction. + +## Python interface + +A task policy must be callable, or define `predict_chunk(observation)`, and +return joint actions with shape `[time, 14]`. Outputs shaped `[1, time, 14]`, +PyTorch tensors, and dictionaries containing `actions`, `action`, or `chunk` are +normalized automatically. + +```python +from policy_speed_env import create_speed_env +from speed_policy import FixedSpeedPolicy, rollout_speed_policy + + +class MyChunkPolicy: + def reset(self): + pass + + def predict_chunk(self, observation): + return model(observation) + + +env = create_speed_env( + "insertion", + chunk_predictor=MyChunkPolicy(), + seed=0, +) +result = rollout_speed_policy(env, FixedSpeedPolicy(1.5)) +``` + +At each speed-policy decision, the environment requests a fresh receding-horizon +chunk, interpolates it at the selected speed, and holds that speed for the +configured frame-skip block. A short chunk is safely replanned as needed. + +`TorchChunkPredictor` in `chunked_policy.py` provides normalization and tensor +handling for ACT-style models. + +## CLI factory + +Expose a factory in the external policy package: + +```python +# my_policy/integration.py +def build_policy(task_name, checkpoint, device): + return MyChunkPolicy.load(checkpoint, task=task_name, device=device) +``` + +Install that package in the SpeedTuning environment, then reference the factory +as `module:attribute`: + +```bash +uv run speedtuning-eval-speed \ + --task insertion \ + --base-policy external-chunk \ + --chunk-policy my_policy.integration:build_policy \ + --upstream-checkpoint /path/to/upstream.pt \ + --speed-policy fixed --speed 1.5 \ + --episodes 10 +``` + +The factory may accept any subset of `task_name`, `checkpoint`, `device`, and +keys supplied through `--factory-kwargs`. + +Training uses the same task-policy flags: + +```bash +uv run speedtuning-train-speed \ + --task insertion \ + --base-policy external-chunk \ + --chunk-policy my_policy.integration:build_policy \ + --upstream-checkpoint /path/to/upstream.pt \ + --output outputs/insertion_speed.pt +``` + +## Retained ACT integration + +Install the learned-policy dependencies: + +```bash +uv sync --extra rl --extra learned +``` + +The retained ACT loader accepts checkpoints containing `model_state_dict`, +`policy_config`, and normalization arrays under `stats`. The required statistics +are `qpos_mean`, `qpos_std`, `action_mean`, and `action_std`. They may instead be +provided in an `.npz`, JSON, or legacy pickle file. + +```bash +uv run speedtuning-eval-speed \ + --task insertion \ + --base-policy external-chunk \ + --chunk-policy act_integration:build_act_chunk_predictor \ + --upstream-checkpoint /path/to/act.pt \ + --factory-kwargs '{"stats_path":"/path/to/dataset_stats.npz"}' \ + --speed-policy fixed --speed 1.5 \ + --episodes 10 +``` + +ACT checkpoints and legacy pickle statistics require Python pickle loading. Only +load files created locally or obtained from a trusted source. Rainbow speed +checkpoints produced by this repository use PyTorch's restricted weight loader. + +## Speed-policy observations + +The speed learner supports: + +- `--speed-observation state` for selected simulator and proprioceptive fields; +- `--speed-observation visual` for joint state plus encoded camera images; +- `--speed-observation external` for a custom feature encoder. + +Use `--no-env-state` to remove privileged simulator state. Visual runs support +pretrained or randomly initialized ResNet-18 encoders, configurable cameras, and +frame stacking. Encoder state and preprocessing metadata are saved in the local +Rainbow checkpoint so evaluation cannot silently change them. + +The retained ACT backbone can be used through: + +```text +--speed-observation external \ +--observation-encoder-loader act_integration:build_act_observation_encoder +``` + +## Archival paper configuration + +The `paper-sim` preset records the visual observation and speed-action settings +recoverable from the paper and retained experiment names. It requires an +external task-policy checkpoint and is provided as an archival starting point, +not an exact numerical reproduction claim. + +```bash +uv run speedtuning-train-speed \ + --config paper-sim \ + --task tea_bag \ + --chunk-policy my_policy.integration:build_policy \ + --upstream-checkpoint /path/to/task_policy.pt \ + --output outputs/tea_bag_speed.pt +``` + +The JSON manifests under `configs/ablations/` cover the observation, reward, +frame-skip, image-encoder, and task-policy ablations retained for research use. diff --git a/docs/SCRIPTED_REPRODUCTION.md b/docs/SCRIPTED_REPRODUCTION.md new file mode 100644 index 0000000..cd172c1 --- /dev/null +++ b/docs/SCRIPTED_REPRODUCTION.md @@ -0,0 +1,154 @@ +# Scripted-policy simulation reproduction + +This guide reproduces the simulation methodology from *SpeedTuning: Speeding Up +Policy Execution with Lightweight Reinforcement Learning* using the three task +policies included in the repository. + +No pretrained model is required. Rainbow DQN learns the execution-speed policy +from scratch, and all generated checkpoints and reports stay under the ignored +`outputs/` directory. + +## Method + +The public presets use the retained scripted-policy training recipe: + +- state observations containing object state, joint position, and joint velocity; +- one speed decision every 10 MuJoCo physics steps; +- reward `0.01 * speed**2` per physics step and 100 for terminal success; +- categorical dueling Double DQN with NoisyNet, prioritized replay, and 3-step + returns; +- one optimizer update per speed decision, applied at episode boundaries; +- 100,000 speed decisions with periodic local checkpoints. + +Task-specific presets expose appropriate discrete speed ranges: + +| Task | Preset | Speed range | +| --- | --- | --- | +| Pick-and-place | `scripted-pick-and-place` | 1.0x-4.5x | +| Insertion | `scripted-insertion` | 1.0x-3.0x | +| Tea bag | `scripted-tea-bag` | 1.0x-3.0x | +| Tea bag, randomized poses | `scripted-tea-bag-randomized` | 1.0x-3.0x | + +The complete hyperparameters are versioned in `configs/`. Command-line options +can override them for ablations or short smoke runs. + +## Install and verify + +```bash +uv sync --extra rl --extra test +MUJOCO_GL=egl uv run pytest -q +MUJOCO_GL=egl uv run speedtuning-sim +``` + +`MUJOCO_GL=egl` enables headless rendering on Linux. Depending on the platform, +DM Control may select a suitable backend without it. + +## Train + +Run each task separately because every task has its own speed policy: + +```bash +MUJOCO_GL=egl uv run speedtuning-train-speed \ + --config scripted-pick-and-place --task pick_and_place \ + --output outputs/pick_and_place_speed.pt \ + --report outputs/pick_and_place_speed.training.json --quiet + +MUJOCO_GL=egl uv run speedtuning-train-speed \ + --config scripted-insertion --task insertion \ + --output outputs/insertion_speed.pt \ + --report outputs/insertion_speed.training.json --quiet + +MUJOCO_GL=egl uv run speedtuning-train-speed \ + --config scripted-tea-bag --task tea_bag \ + --output outputs/tea_bag_speed.pt \ + --report outputs/tea_bag_speed.training.json --quiet +``` + +Training defaults to CPU. Add `--device cuda` on a CUDA machine. Hardware changes +runtime, but it does not change the simulator protocol or reported acceleration +metric. + +The presets write a snapshot every 10,000 decisions. Checkpoint metadata records +the task, pose protocol, speed actions, observation preprocessing, and training +configuration. Evaluation rejects an incompatible task or protocol. + +## Evaluate + +Evaluate held-out initial poses by choosing a starting seed and episode count: + +```bash +MUJOCO_GL=egl uv run speedtuning-eval-speed \ + --config scripted-pick-and-place --task pick_and_place \ + --speed-policy rainbow \ + --speed-checkpoint outputs/pick_and_place_speed.pt \ + --seed 100 --episodes 100 +``` + +Compare against a fixed speed on the same seeds: + +```bash +MUJOCO_GL=egl uv run speedtuning-eval-speed \ + --config scripted-pick-and-place --task pick_and_place \ + --speed-policy fixed --speed 3.856 \ + --seed 100 --episodes 100 +``` + +For a complete fixed-speed frontier: + +```bash +MUJOCO_GL=egl uv run speedtuning-sweep \ + --config scripted-pick-and-place --task pick_and_place \ + --speed-start 1.0 --speed-stop 4.5 --speed-step 0.25 \ + --episodes-per-speed 100 \ + --output outputs/pick_and_place_sweep.json +``` + +Physical acceleration is the nominal task horizon divided by executed MuJoCo +steps. Failed rollouts remain in the success-rate and acceleration summaries. + +## Tea-bag protocols + +The retained tea-bag environment has one fixed initial object pose. It is useful +for matching the historical scripted-policy experiment, but repeated seeds are +intentionally identical. + +Use `scripted-tea-bag-randomized` when a distribution of initial poses is needed: + +```bash +MUJOCO_GL=egl uv run speedtuning-train-speed \ + --config scripted-tea-bag-randomized --task tea_bag --seed 1 \ + --output outputs/tea_bag_randomized_speed.pt --quiet + +MUJOCO_GL=egl uv run speedtuning-eval-speed \ + --config scripted-tea-bag-randomized --task tea_bag \ + --speed-policy rainbow \ + --speed-checkpoint outputs/tea_bag_randomized_speed.pt \ + --seed 100 --episodes 100 +``` + +Do not evaluate a fixed-pose checkpoint under the randomized protocol or vice +versa; the observation distribution and normalization differ. + +## Reference results + +One seeded run using the protocol above produced: + +| Protocol | Learned speed | Matched fixed speed | +| --- | --- | --- | +| Pick-and-place, seeds 100-199 | 98% at 3.856x | 66% at 3.846x | +| Insertion, seeds 100-199 | 97% at 2.387x | 52% at 2.381x | +| Tea bag, randomized seeds 100-199 | 78% at 2.077x | 24% at 2.075x | + +The machine-readable record is +[`benchmarks/scripted_results.json`](../benchmarks/scripted_results.json). +Reinforcement learning is stochastic, so compare reruns at the level of success +and acceleration trends rather than exact decimal equality. + +## Scope + +This workflow reproduces speed-policy learning, temporal acceleration, sparse +success, fixed-speed baselines, and seeded simulation evaluation. It does not +include real-robot execution, datasets, or learned task-policy checkpoints. + +See [External task-policy integration](EXTERNAL_POLICIES.md) to wrap ACT or +another action-chunk policy. diff --git a/docs/assets/speedtuning_teaser.png b/docs/assets/speedtuning_teaser.png new file mode 100644 index 0000000..bee7b59 Binary files /dev/null and b/docs/assets/speedtuning_teaser.png differ diff --git a/ee_sim_env.py b/ee_sim_env.py new file mode 100644 index 0000000..8c145af --- /dev/null +++ b/ee_sim_env.py @@ -0,0 +1,235 @@ +"""MuJoCo environments controlled by left and right end-effector poses.""" + +from __future__ import annotations + +import collections +import os + +import numpy as np +from dm_control import mujoco +from dm_control.rl import control +from dm_control.suite import base + +from constants import ( + DT, + PUPPET_GRIPPER_POSITION_CLOSE, + PUPPET_GRIPPER_POSITION_NORMALIZE_FN, + PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN, + PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN, + START_ARM_POSE, + XML_DIR, +) +from sim_tasks import ( + get_task_spec, + insertion_reward, + normalize_task_name, + sample_box_pose, + sample_insertion_pose, + tea_bag_reward, + transfer_cube_reward, +) + + +# Kept for compatibility with the historical scripts. New code should pass +# render_images=False to make_ee_sim_env instead. +DISABLE_RENDER = [False] + + +def make_ee_sim_env( + task_name: str, + render_images: bool = True, + seed: int | None = None, + object_pose=None, + randomize_object_pose: bool = False, +): + """Create an end-effector-control environment for a SpeedTuning sim task. + + Actions contain two 8D commands: xyz, quaternion, and normalized gripper + position for the left arm, followed by the same fields for the right arm. + The factory accepts both the public task names (``pick_and_place``, + ``insertion``, ``tea_bag``) and the task names used by the historical code. + """ + + task_name = normalize_task_name(task_name) + spec = get_task_spec(task_name) + physics = mujoco.Physics.from_xml_path(os.path.join(XML_DIR, spec.ee_xml)) + random_state = np.random.RandomState(seed) + task_classes = { + "pick_and_place": TransferCubeEETask, + "insertion": InsertionEETask, + "tea_bag": TransferTeaBagEETask, + } + task = task_classes[task_name]( + random=random_state, + render_images=render_images, + object_pose=object_pose, + randomize_object_pose=randomize_object_pose, + ) + return control.Environment( + physics, + task, + time_limit=20, + control_timestep=DT, + n_sub_steps=None, + flat_observation=False, + ) + + +class BimanualViperXEETask(base.Task): + def __init__( + self, + random=None, + render_images: bool = True, + object_pose=None, + randomize_object_pose: bool = False, + ): + super().__init__(random=random) + self.render_images = render_images + self.object_pose = ( + None if object_pose is None else np.asarray(object_pose).copy() + ) + self.randomize_object_pose = bool(randomize_object_pose) + + def before_step(self, action, physics): + action = np.asarray(action, dtype=np.float64) + if action.shape != (16,) or not np.all(np.isfinite(action)): + raise ValueError("End-effector actions must be a finite array with shape (16,)") + + action_left = action[:8] + action_right = action[8:] + np.copyto(physics.data.mocap_pos[0], action_left[:3]) + np.copyto(physics.data.mocap_quat[0], action_left[3:7]) + np.copyto(physics.data.mocap_pos[1], action_right[:3]) + np.copyto(physics.data.mocap_quat[1], action_right[3:7]) + + left_gripper = PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN(action_left[7]) + right_gripper = PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN(action_right[7]) + np.copyto( + physics.data.ctrl, + [left_gripper, -left_gripper, right_gripper, -right_gripper], + ) + + def initialize_robots(self, physics): + physics.named.data.qpos[:16] = START_ARM_POSE + np.copyto(physics.data.mocap_pos[0], [-0.31718881, 0.5, 0.29525084]) + np.copyto(physics.data.mocap_quat[0], [1, 0, 0, 0]) + np.copyto(physics.data.mocap_pos[1], [0.31718881, 0.49999888, 0.29525084]) + np.copyto(physics.data.mocap_quat[1], [1, 0, 0, 0]) + np.copyto( + physics.data.ctrl, + [ + PUPPET_GRIPPER_POSITION_CLOSE, + -PUPPET_GRIPPER_POSITION_CLOSE, + PUPPET_GRIPPER_POSITION_CLOSE, + -PUPPET_GRIPPER_POSITION_CLOSE, + ], + ) + + @staticmethod + def get_qpos(physics): + qpos = physics.data.qpos.copy() + left, right = qpos[:8], qpos[8:16] + return np.concatenate( + [ + left[:6], + [PUPPET_GRIPPER_POSITION_NORMALIZE_FN(left[6])], + right[:6], + [PUPPET_GRIPPER_POSITION_NORMALIZE_FN(right[6])], + ] + ) + + @staticmethod + def get_qvel(physics): + qvel = physics.data.qvel.copy() + left, right = qvel[:8], qvel[8:16] + return np.concatenate( + [ + left[:6], + [PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN(left[6])], + right[:6], + [PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN(right[6])], + ] + ) + + @staticmethod + def get_env_state(physics): + return physics.data.qpos.copy()[16:] + + def get_observation(self, physics): + obs = collections.OrderedDict( + qpos=self.get_qpos(physics), + qvel=self.get_qvel(physics), + env_state=self.get_env_state(physics), + ) + if self.render_images and not DISABLE_RENDER[0]: + obs["images"] = { + "top": physics.render(height=480, width=640, camera_id="top"), + "angle": physics.render(height=480, width=640, camera_id="angle"), + "vis": physics.render(height=480, width=640, camera_id="front_close"), + } + obs["mocap_pose_left"] = np.concatenate( + [physics.data.mocap_pos[0], physics.data.mocap_quat[0]] + ).copy() + obs["mocap_pose_right"] = np.concatenate( + [physics.data.mocap_pos[1], physics.data.mocap_quat[1]] + ).copy() + obs["gripper_ctrl"] = physics.data.ctrl.copy() + return obs + + +class TransferCubeEETask(BimanualViperXEETask): + max_reward = 4 + + def initialize_episode(self, physics): + self.initialize_robots(physics) + pose = ( + sample_box_pose(self.random) + if self.object_pose is None + else self.object_pose + ) + if np.asarray(pose).shape != (7,): + raise ValueError("Pick-and-place object_pose must have shape (7,)") + physics.named.data.qpos["red_box_joint"] = pose + super().initialize_episode(physics) + + def get_reward(self, physics): + return transfer_cube_reward(physics) + + +class InsertionEETask(BimanualViperXEETask): + max_reward = 4 + + def initialize_episode(self, physics): + self.initialize_robots(physics) + if self.object_pose is None: + peg_pose, socket_pose = sample_insertion_pose(self.random) + else: + if self.object_pose.shape != (14,): + raise ValueError("Insertion object_pose must have shape (14,)") + peg_pose, socket_pose = self.object_pose[:7], self.object_pose[7:] + physics.named.data.qpos["red_peg_joint"] = peg_pose + physics.named.data.qpos["blue_socket_joint"] = socket_pose + super().initialize_episode(physics) + + def get_reward(self, physics): + return insertion_reward(physics) + + +class TransferTeaBagEETask(BimanualViperXEETask): + max_reward = 3 + + def initialize_episode(self, physics): + self.initialize_robots(physics) + if self.object_pose is not None: + pose = self.object_pose + elif self.randomize_object_pose: + pose = sample_box_pose(self.random) + else: + pose = [0.15, 0.5, 0.05, 1, 0, 0, 0] + if np.asarray(pose).shape != (7,): + raise ValueError("Tea-bag object_pose must have shape (7,)") + physics.named.data.qpos["red_box_joint"] = pose + super().initialize_episode(physics) + + def get_reward(self, physics): + return tea_bag_reward(physics) diff --git a/experiment_config.py b/experiment_config.py new file mode 100644 index 0000000..14bf749 --- /dev/null +++ b/experiment_config.py @@ -0,0 +1,73 @@ +"""JSON experiment manifests shared by training, evaluation, and sweeps.""" + +from __future__ import annotations + +import argparse +import json +from importlib import resources +from pathlib import Path + + +CONFIG_ROOT = Path(str(resources.files("configs"))) +NAMED_CONFIGS = { + "paper-sim": CONFIG_ROOT / "paper_sim.json", + "scripted-tea-bag": CONFIG_ROOT / "scripted_tea_bag.json", + "scripted-tea-bag-randomized": ( + CONFIG_ROOT / "scripted_tea_bag_randomized.json" + ), + "scripted-pick-and-place": ( + CONFIG_ROOT / "scripted_pick_and_place.json" + ), + "scripted-insertion": CONFIG_ROOT / "scripted_insertion.json", +} +CONFIG_SECTIONS = ( + "base_policy", + "environment", + "observation", + "reward", + "training", + "evaluation", +) + + +def resolve_config_path(value): + if value is None: + return None + if value in NAMED_CONFIGS: + return NAMED_CONFIGS[value] + return Path(value) + + +def load_experiment_config(value): + path = resolve_config_path(value) + if path is None: + return {}, None + if not path.exists(): + raise ValueError(f"Experiment config does not exist: {path}") + payload = json.loads(path.read_text()) + if not isinstance(payload, dict): + raise ValueError("Experiment config must contain a JSON object") + inherited_defaults = {} + inherited_manifest = None + if payload.get("inherits"): + inherited_defaults, inherited_manifest = load_experiment_config( + payload["inherits"] + ) + defaults = dict(inherited_defaults) + for section in CONFIG_SECTIONS: + values = payload.get(section, {}) + if not isinstance(values, dict): + raise ValueError(f"Experiment config section {section!r} must be an object") + defaults.update(values) + return defaults, { + "path": str(path), + "manifest": payload, + "inherits": inherited_manifest, + } + + +def defaults_from_argv(argv=None): + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--config") + known, _ = parser.parse_known_args(argv) + return load_experiment_config(known.config) diff --git a/policy.py b/policy.py new file mode 100644 index 0000000..accb743 --- /dev/null +++ b/policy.py @@ -0,0 +1,81 @@ +import torch.nn as nn +from torch.nn import functional as F +import torchvision.transforms as transforms + +from detr.main import build_ACT_model_and_optimizer, build_CNNMLP_model_and_optimizer + +class ACTPolicy(nn.Module): + def __init__(self, args_override): + super().__init__() + model, optimizer = build_ACT_model_and_optimizer(args_override) + self.model = model # CVAE decoder + self.optimizer = optimizer + self.kl_weight = args_override['kl_weight'] + + def __call__(self, qpos, image, actions=None, is_pad=None): + env_state = None + normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225]) + image = normalize(image) + if actions is not None: # training time + actions = actions[:, :self.model.num_queries] + is_pad = is_pad[:, :self.model.num_queries] + + a_hat, is_pad_hat, (mu, logvar) = self.model(qpos, image, env_state, actions, is_pad) + total_kld, dim_wise_kld, mean_kld = kl_divergence(mu, logvar) + loss_dict = dict() + all_l1 = F.l1_loss(actions, a_hat, reduction='none') + l1 = (all_l1 * ~is_pad.unsqueeze(-1)).mean() + loss_dict['l1'] = l1 + loss_dict['kl'] = total_kld[0] + loss_dict['loss'] = loss_dict['l1'] + loss_dict['kl'] * self.kl_weight + return loss_dict + else: # inference time + a_hat, _, (_, _) = self.model(qpos, image, env_state) # no action, sample from prior + return a_hat + + def configure_optimizers(self): + return self.optimizer + + +class CNNMLPPolicy(nn.Module): + def __init__(self, args_override): + super().__init__() + model, optimizer = build_CNNMLP_model_and_optimizer(args_override) + self.model = model # decoder + self.optimizer = optimizer + + def __call__(self, qpos, image, actions=None, is_pad=None): + env_state = None # TODO + normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225]) + image = normalize(image) + if actions is not None: # training time + actions = actions[:, 0] + a_hat = self.model(qpos, image, env_state, actions) + mse = F.mse_loss(actions, a_hat) + loss_dict = dict() + loss_dict['mse'] = mse + loss_dict['loss'] = loss_dict['mse'] + return loss_dict + else: # inference time + a_hat = self.model(qpos, image, env_state) # no action, sample from prior + return a_hat + + def configure_optimizers(self): + return self.optimizer + +def kl_divergence(mu, logvar): + batch_size = mu.size(0) + assert batch_size != 0 + if mu.data.ndimension() == 4: + mu = mu.view(mu.size(0), mu.size(1)) + if logvar.data.ndimension() == 4: + logvar = logvar.view(logvar.size(0), logvar.size(1)) + + klds = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()) + total_kld = klds.sum(1).mean(0, True) + dimension_wise_kld = klds.mean(0) + mean_kld = klds.mean(1).mean(0, True) + + return total_kld, dimension_wise_kld, mean_kld diff --git a/policy_loader.py b/policy_loader.py new file mode 100644 index 0000000..68df88a --- /dev/null +++ b/policy_loader.py @@ -0,0 +1,97 @@ +"""Small entry-point loader for integrating external policy repositories.""" + +from __future__ import annotations + +import importlib +import inspect +from pathlib import Path + +from chunked_policy import ChunkPredictorAdapter +from speed_policy import SpeedPolicyAdapter +from speed_observation import ObservationEncoderAdapter + + +def load_entrypoint(spec: str): + """Load ``module.submodule:attribute`` without modifying ``sys.path``.""" + + if ":" not in spec: + raise ValueError("An entry point must use the form 'module.submodule:attribute'") + module_name, attribute_name = spec.rsplit(":", 1) + if not module_name or not attribute_name: + raise ValueError("An entry point must use the form 'module.submodule:attribute'") + module = importlib.import_module(module_name) + try: + return getattr(module, attribute_name) + except AttributeError as exc: + raise ValueError(f"{module_name!r} has no attribute {attribute_name!r}") from exc + + +def _call_factory(factory, available_kwargs): + if not callable(factory): + return factory + signature = inspect.signature(factory) + accepts_extra = any( + parameter.kind == inspect.Parameter.VAR_KEYWORD + for parameter in signature.parameters.values() + ) + kwargs = { + name: value + for name, value in available_kwargs.items() + if accepts_extra or name in signature.parameters + } + return factory(**kwargs) + + +def load_chunk_predictor( + entrypoint, + task_name, + checkpoint=None, + device="cpu", + factory_kwargs=None, +): + """Instantiate and validate an external joint-action chunk predictor. + + The factory may accept any subset of ``task_name``, ``checkpoint``, + ``device``, and the keys in ``factory_kwargs``. Its result must be callable + or define ``predict_chunk(observation)``. + """ + + factory = load_entrypoint(entrypoint) + available = { + "task_name": task_name, + "checkpoint": None if checkpoint is None else Path(checkpoint), + "device": device, + **(factory_kwargs or {}), + } + return ChunkPredictorAdapter(_call_factory(factory, available)) + + +def load_speed_policy(entrypoint, checkpoint=None, device="cpu", factory_kwargs=None): + """Instantiate and validate an external physical-speed policy.""" + + factory = load_entrypoint(entrypoint) + available = { + "checkpoint": None if checkpoint is None else Path(checkpoint), + "device": device, + **(factory_kwargs or {}), + } + return SpeedPolicyAdapter(_call_factory(factory, available)) + + +def load_observation_encoder( + entrypoint, + task_name, + checkpoint=None, + device="cpu", + factory_kwargs=None, +): + """Instantiate an external speed-observation encoder.""" + + factory = load_entrypoint(entrypoint) + available = { + "task_name": task_name, + "checkpoint": None if checkpoint is None else Path(checkpoint), + "device": device, + **(factory_kwargs or {}), + } + return ObservationEncoderAdapter(_call_factory(factory, available)) diff --git a/policy_speed_env.py b/policy_speed_env.py new file mode 100644 index 0000000..c4e527a --- /dev/null +++ b/policy_speed_env.py @@ -0,0 +1,515 @@ +"""Speed-control environments for scripted and chunked robot policies. + +The speed agent never produces robot actions. It observes simulator state and +chooses how quickly an independent base policy advances through nominal policy +time. This separation lets the same speed learner wrap retained waypoint +controllers, ACT, or another upstream policy that emits 14D action chunks. +""" + +from __future__ import annotations + +from collections import deque +from pathlib import Path +from typing import Callable, Sequence + +import numpy as np +from dm_control.rl import control + +from chunked_policy import ( + ChunkedPolicyRunner, + RecordedChunkPredictor, + collect_scripted_joint_demonstration, +) +from ee_sim_env import make_ee_sim_env +from scripted_policy import make_scripted_policy +from sim_env import make_sim_env +from sim_tasks import get_task_spec, normalize_task_name +from speed_observation import StateObservationEncoder, encoder_spec + + +DEFAULT_SPEEDS = (1.0, 1.5, 2.0, 2.5, 3.0) + + +def terminal_success_reward(speed, done, success): + """Sparse objective used by the initial SpeedTuning experiments.""" + + del speed + return 100.0 if done and success else 0.0 + + +def make_speed_reward(success_bonus=100.0, speed_weight=0.01, speed_power=2.0): + """Build a serializable-style speed objective for public experiments.""" + + if speed_power < 0: + raise ValueError("speed_power must be non-negative") + + def reward_fn(speed, done, success): + reward = speed_weight * float(speed) ** speed_power + if done and success: + reward += success_bonus + return reward + + return reward_fn + + +def encode_state_observation(observation): + """Flatten task state, robot position, and robot velocity for a speed agent.""" + + return StateObservationEncoder()(observation) + + +class WaypointActionSource: + """Advance an end-effector waypoint controller by a requested speed.""" + + def __init__(self, policy): + self.policy = policy + + def reset(self): + reset = getattr(self.policy, "reset", None) + if reset is not None: + reset() + + def begin_decision(self, timestep, speed): + del timestep, speed + + def action(self, timestep, speed): + return self.policy(timestep, step_inc=speed) + + +class ChunkedActionSource: + """Advance an arbitrary joint-action chunk predictor by a requested speed.""" + + def __init__(self, predictor): + self.runner = ChunkedPolicyRunner(predictor) + + def reset(self): + self.runner.reset() + + def begin_decision(self, timestep, speed): + self.runner.begin_decision(timestep.observation, speed=speed) + + def action(self, timestep, speed): + return self.runner.action(timestep.observation, speed=speed) + + +class SpeedPolicyEnv: + """A small RL environment whose actions select base-policy speed. + + The interface follows the classic ``reset() -> observation`` and + ``step(action) -> observation, reward, done, info`` convention used by the + included Rainbow implementation. Discrete action ``i`` maps to + ``speed_values[i]``; callers may also pass a physical speed with + ``quantized=False``. + """ + + def __init__( + self, + env, + action_source, + episode_len, + reward_fn=None, + speed_values: Sequence[float] = DEFAULT_SPEEDS, + observation_encoder: Callable[[dict], np.ndarray] | None = None, + frame_stack=1, + decision_frame_skip=10, + save_video=False, + onscreen_render=False, + video_path="output_video.mp4", + max_physics_steps=None, + terminate_on_success=False, + environment_metadata=None, + ): + self.env = env + self.action_source = action_source + self.reward_fn = reward_fn or terminal_success_reward + self.episode_len = int(episode_len) + self.observation_encoder = observation_encoder or StateObservationEncoder() + self.frame_stack = int(frame_stack) + self.decision_frame_skip = int(decision_frame_skip) + if self.frame_stack <= 0 or self.decision_frame_skip <= 0: + raise ValueError("frame_stack and decision_frame_skip must be positive") + self.save_video = bool(save_video) + self.onscreen_render = bool(onscreen_render) + self.video_path = Path(video_path) + self.terminate_on_success = bool(terminate_on_success) + self._environment_metadata = dict(environment_metadata or {}) + + values = np.asarray(speed_values, dtype=np.float64) + if values.ndim != 1 or len(values) == 0: + raise ValueError("speed_values must be a non-empty one-dimensional sequence") + if not np.all(np.isfinite(values)) or np.any(values <= 0): + raise ValueError("Every speed value must be finite and positive") + self.speed_values = tuple(float(value) for value in values) + self.action_space = len(self.speed_values) + self.max_physics_steps = int( + max_physics_steps + if max_physics_steps is not None + else np.ceil(self.episode_len / min(self.speed_values)) + ) + if self.max_physics_steps <= 0: + raise ValueError("max_physics_steps must be positive") + + env_state_dim = int(self.env.physics.data.qpos[16:].size) + output_dim = getattr(self.observation_encoder, "output_dim", None) + try: + single_observation_dim = int(output_dim(env_state_dim)) + except (AttributeError, TypeError): + single_observation_dim = env_state_dim + 28 + self.obs_space = single_observation_dim * self.frame_stack + self.cur_ts = None + self.cur_success = False + self.policy_time = 0.0 + self.physics_steps = 0 + self.speed_list = [] + self.image_list = [] + self._observation_stack = deque(maxlen=self.frame_stack) + self._figure = None + self._plot_image = None + + def reset(self): + self.cur_ts = self.env.reset() + self.action_source.reset() + reset_encoder = getattr(self.observation_encoder, "reset", None) + if reset_encoder is not None: + reset_encoder() + self.cur_success = False + self.policy_time = 0.0 + self.physics_steps = 0 + self.speed_list = [] + self.image_list = [] + self._observation_stack.clear() + + if self.onscreen_render: + import matplotlib.pyplot as plt + + self._figure, axis = plt.subplots() + self._plot_image = axis.imshow(self.cur_ts.observation["images"]["angle"]) + plt.ion() + encoded = self._encode_observation(self.cur_ts.observation) + for _ in range(self.frame_stack): + self._observation_stack.append(encoded.copy()) + observation = self.get_obs() + self.obs_space = int(observation.size) + return observation + + def speed_from_action(self, action): + if not isinstance(action, (int, np.integer)): + raise ValueError("A discrete speed action must be an integer index") + action = int(action) + if not 0 <= action < self.action_space: + raise ValueError(f"Speed action must be in [0, {self.action_space - 1}]") + return self.speed_values[action] + + @staticmethod + def _validate_continuous_speed(speed): + speed = float(speed) + if not np.isfinite(speed) or speed <= 0: + raise ValueError("Continuous speed must be finite and positive") + return speed + + def _resolve_speed(self, speed, quantized): + return ( + self.speed_from_action(speed) + if quantized + else self._validate_continuous_speed(speed) + ) + + def begin_decision(self, speed, quantized=True): + """Select one speed and force a fresh receding-horizon task chunk.""" + + if self.cur_ts is None: + raise RuntimeError("Call reset() before beginning a decision") + resolved_speed = self._resolve_speed(speed, quantized) + begin = getattr(self.action_source, "begin_decision", None) + if begin is not None: + begin(self.cur_ts, resolved_speed) + return resolved_speed + + def step(self, speed, quantized=True): + """Advance one physics step, retaining legacy low-level behavior.""" + + return self._step_physics(self._resolve_speed(speed, quantized)) + + def step_decision(self, speed, frame_skip=None, quantized=True): + """Execute one speed-policy action for a fixed block of physics steps. + + A fresh task-policy chunk is predicted at the decision boundary. Rewards + inside the block are summed without discounting and one transition is + returned to the speed learner. + """ + + repeats = self.decision_frame_skip if frame_skip is None else int(frame_skip) + if repeats <= 0: + raise ValueError("frame_skip must be positive") + resolved_speed = self.begin_decision(speed, quantized=quantized) + total_reward = 0.0 + executed = 0 + info = {"success": False} + for _ in range(repeats): + observation, reward, done, info = self._step_physics(resolved_speed) + total_reward += float(reward) + executed += 1 + if done: + break + info = dict(info) + info.update( + decision_frame_skip=repeats, + decision_physics_steps=executed, + reward_aggregation="undiscounted_sum", + ) + return observation, total_reward, done, info + + def _step_physics(self, speed): + if self.cur_ts is None: + raise RuntimeError("Call reset() before step()") + if self.policy_time >= self.episode_len or self.physics_steps >= self.max_physics_steps: + raise RuntimeError("The episode has already finished") + + action = self.action_source.action(self.cur_ts, speed) + try: + next_timestep = self.env.step(action) + except control.PhysicsError as exc: + # Aggressive speed exploration can occasionally drive MuJoCo into an + # invalid state. Treat that rollout as a failed terminal transition + # so one unstable sample cannot abort a long reinforcement-learning + # run. The following reset restores the simulator state. + self.policy_time += speed + self.physics_steps += 1 + self.speed_list.append(speed) + done = True + reward = float(self.reward_fn(speed, done, False)) + info = { + "success": False, + "speed": speed, + "policy_time": self.policy_time, + "physics_steps": self.physics_steps, + "task_reward": 0.0, + "target_reward": self.env.task.max_reward, + "physics_error": str(exc), + } + return self.get_obs(), reward, done, info + self.cur_ts = next_timestep + self.policy_time += speed + self.physics_steps += 1 + self.speed_list.append(speed) + self._observation_stack.append( + self._encode_observation(self.cur_ts.observation) + ) + + if self.onscreen_render: + import matplotlib.pyplot as plt + + self._plot_image.set_data(self.cur_ts.observation["images"]["angle"]) + plt.pause(0.002) + if self.save_video: + self.image_list.append(self.cur_ts.observation["images"]["angle"]) + + task_reward = float(self.cur_ts.reward or 0) + if task_reward >= self.env.task.max_reward: + self.cur_success = True + timed_out = ( + self.policy_time >= self.episode_len + or self.physics_steps >= self.max_physics_steps + ) + done = timed_out or (self.terminate_on_success and self.cur_success) + reward = float(self.reward_fn(speed, done, self.cur_success)) + + if done and self.save_video: + self._save_video() + info = { + "success": self.cur_success, + "speed": speed, + "policy_time": self.policy_time, + "physics_steps": self.physics_steps, + "task_reward": task_reward, + "target_reward": self.env.task.max_reward, + } + return self.get_obs(), reward, done, info + + def _encode_observation(self, observation_dict): + observation = np.asarray( + self.observation_encoder(observation_dict), dtype=np.float32 + ) + if observation.ndim != 1 or not np.all(np.isfinite(observation)): + raise ValueError("The speed-policy observation must be a finite 1D array") + return observation + + def get_obs(self): + if self.cur_ts is None or not self._observation_stack: + raise RuntimeError("Call reset() before requesting an observation") + return np.concatenate(tuple(self._observation_stack)).astype( + np.float32, copy=False + ) + + def observation_spec(self): + return { + "encoder": encoder_spec(self.observation_encoder), + "frame_stack": self.frame_stack, + "padding": "repeat_initial", + "observation_dim": int(self.obs_space), + } + + def environment_spec(self): + return dict(self._environment_metadata) + + def observation_encoder_state_dict(self): + state_dict = getattr(self.observation_encoder, "state_dict", None) + return None if state_dict is None else state_dict() + + def load_observation_encoder_state_dict(self, state_dict): + if state_dict is None: + return + load = getattr(self.observation_encoder, "load_state_dict", None) + if load is None: + raise ValueError("Configured observation encoder cannot load checkpoint state") + load(state_dict) + + def _save_video(self): + try: + import imageio.v2 as imageio + except ImportError as exc: + raise RuntimeError("Video export requires: uv sync --extra video") from exc + self.video_path.parent.mkdir(parents=True, exist_ok=True) + imageio.mimsave(self.video_path, self.image_list, fps=50) + + def close(self): + if self._figure is not None: + import matplotlib.pyplot as plt + + plt.close(self._figure) + self._figure = None + close_encoder = getattr(self.observation_encoder, "close", None) + if close_encoder is not None: + close_encoder() + + +def create_speed_env( + task_name="tea_bag", + reward_fn=None, + chunk_predictor=None, + object_pose=None, + onscreen_render=False, + save_video=False, + render_images=None, + seed=None, + speed_values=DEFAULT_SPEEDS, + video_path="output_video.mp4", + observation_encoder=None, + frame_stack=1, + decision_frame_skip=10, + terminate_on_success=False, + randomize_object_pose=False, +): + """Create a speed environment around a scripted or chunked base policy. + + With no ``chunk_predictor`` this uses the retained end-effector waypoint + controller. Supplying any callable chunk predictor switches to the joint + simulator and the generic ``[time, 14]`` action-chunk contract. + """ + + task_name = normalize_task_name(task_name) + spec = get_task_spec(task_name) + observation_encoder = observation_encoder or StateObservationEncoder() + if render_images is None: + render_images = bool( + onscreen_render + or save_video + or chunk_predictor is not None + or getattr(observation_encoder, "requires_images", False) + ) + if getattr(observation_encoder, "requires_images", False) and not render_images: + raise ValueError("The configured speed observation encoder requires images") + + if chunk_predictor is None: + env = make_ee_sim_env( + task_name, + render_images=render_images, + seed=seed, + object_pose=object_pose, + randomize_object_pose=randomize_object_pose, + ) + action_source = WaypointActionSource(make_scripted_policy(task_name)) + else: + env = make_sim_env( + task_name, + render_images=render_images, + seed=seed, + object_pose=object_pose, + randomize_object_pose=randomize_object_pose, + ) + action_source = ChunkedActionSource(chunk_predictor) + + return SpeedPolicyEnv( + env=env, + action_source=action_source, + reward_fn=reward_fn, + episode_len=spec.episode_len, + speed_values=speed_values, + observation_encoder=observation_encoder, + frame_stack=frame_stack, + decision_frame_skip=decision_frame_skip, + onscreen_render=onscreen_render, + save_video=save_video, + video_path=video_path, + terminate_on_success=terminate_on_success, + environment_metadata={ + "task": task_name, + "base_policy": ( + "scripted" if chunk_predictor is None else "chunked" + ), + "randomize_object_pose": bool(randomize_object_pose), + }, + ) + + +def create_recorded_chunk_speed_env( + task_name="tea_bag", + chunk_size=25, + seed=0, + render_images=False, + **kwargs, +): + """Create a checkpoint-free chunk environment for integration testing.""" + + demonstration = collect_scripted_joint_demonstration(task_name, seed=seed) + predictor = RecordedChunkPredictor(demonstration.actions, chunk_size=chunk_size) + return create_speed_env( + task_name=task_name, + chunk_predictor=predictor, + object_pose=demonstration.object_pose, + render_images=render_images, + seed=seed, + **kwargs, + ) + + +def test_speed_env( + task_name="tea_bag", + speed_func=None, + seed=0, + chunk_predictor=None, +): + """Run a continuous speed function and return the maximum wrapper reward.""" + + speed_env = create_speed_env( + task_name=task_name, + chunk_predictor=chunk_predictor, + seed=seed, + ) + observation = speed_env.reset() + done = False + rewards = [] + while not done: + speed = 1.0 if speed_func is None else speed_func( + observation=observation, + policy_time=speed_env.policy_time, + ) + observation, reward, done, _ = speed_env.step(speed, quantized=False) + rewards.append(reward) + speed_env.close() + return max(rewards, default=0.0) + + +if __name__ == "__main__": + for _task_name in ("pick_and_place", "insertion", "tea_bag"): + print(_task_name, test_speed_env(_task_name)) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..e765b8b --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,92 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "speedtuning-sim" +version = "0.1.0" +description = "Simulation tasks and speed-policy infrastructure for SpeedTuning" +readme = "README.md" +requires-python = ">=3.10,<3.11" +license = "MIT" +license-files = ["LICENSE", "NOTICE.md", "detr/LICENSE"] +authors = [{name = "David D. Yuan"}] +keywords = ["robotics", "reinforcement-learning", "mujoco", "action-chunking"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "Operating System :: POSIX :: Linux", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] +dependencies = [ + "dm-control==1.0.9", + "mujoco==2.3.3", + "numpy>=1.23,<2", + "pyquaternion==0.9.9", +] + +[project.optional-dependencies] +learned = [ + "einops>=0.6,<0.9", + "packaging>=23,<26", + "torch==2.5.1", + "torchvision==0.20.1", +] +rl = ["torch==2.5.1"] +test = ["pytest>=8,<9"] +video = ["imageio>=2.31,<3", "imageio-ffmpeg>=0.4.9,<0.7"] +evaluation = ["matplotlib>=3.7,<4"] + +[project.urls] +Paper = "https://doi.org/10.1109/ICRA55743.2025.11128753" +Preprint = "https://arxiv.org/abs/2608.09138" +"Project Page" = "https://daivdyuan.github.io/speed-tuning/" +Video = "https://daivdyuan.github.io/speed-tuning/static/videos/icra2025_final.mp4" +Repository = "https://github.com/DaivdYuan/SpeedTuning" +Documentation = "https://github.com/DaivdYuan/SpeedTuning/blob/main/docs/SCRIPTED_REPRODUCTION.md" +Issues = "https://github.com/DaivdYuan/SpeedTuning/issues" + +[project.scripts] +speedtuning-sim = "scripts.run_sim:main" +speedtuning-check-chunks = "scripts.check_chunked_policy:main" +speedtuning-rainbow-poc = "scripts.rainbow_poc:main" +speedtuning-train-speed = "scripts.train_speed_policy:main" +speedtuning-eval-speed = "scripts.eval_speed_policy:main" +speedtuning-sweep = "scripts.sweep_speed_policy:main" + +[tool.setuptools] +include-package-data = true +py-modules = [ + "act_integration", + "chunked_policy", + "constants", + "ee_sim_env", + "experiment_config", + "policy", + "policy_loader", + "policy_speed_env", + "scripted_policy", + "sim_env", + "sim_tasks", + "speed_policy", + "speed_evaluation", + "speed_observation", + "speed_training", +] + +[tool.setuptools.packages.find] +include = ["assets", "configs*", "detr*", "rl*", "scripts*"] + +[tool.setuptools.package-data] +assets = ["*.xml", "*.stl", "*.obj"] +configs = ["*.json", "*.md", "ablations/*.json"] + +[tool.pytest.ini_options] +pythonpath = ["."] +testpaths = ["tests"] +markers = [ + "learned: requires the optional PyTorch/torchvision learned-policy stack", + "rl: requires the optional PyTorch Rainbow DQN stack", +] diff --git a/requirements-sim.txt b/requirements-sim.txt new file mode 100644 index 0000000..0b48e80 --- /dev/null +++ b/requirements-sim.txt @@ -0,0 +1,6 @@ +# Physics versions are pinned because the historical waypoint controllers are +# sensitive to MuJoCo contact and solver changes. +dm-control==1.0.9 +mujoco==2.3.3 +numpy>=1.23,<2 +pyquaternion==0.9.9 diff --git a/rl/__init__.py b/rl/__init__.py new file mode 100644 index 0000000..273da96 --- /dev/null +++ b/rl/__init__.py @@ -0,0 +1 @@ +"""Reinforcement-learning components for speed selection.""" diff --git a/rl/rainbowDQN/__init__.py b/rl/rainbowDQN/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/rl/rainbowDQN/dqnAgent.py b/rl/rainbowDQN/dqnAgent.py new file mode 100644 index 0000000..dbcc5ee --- /dev/null +++ b/rl/rainbowDQN/dqnAgent.py @@ -0,0 +1,289 @@ +"""Rainbow DQN optimization core used by speed-policy training.""" + +from __future__ import annotations + +import random +from pathlib import Path +from time import perf_counter +from typing import Dict + +import numpy as np +import torch +import torch.optim as optim +from torch.nn.utils import clip_grad_norm_ + +from .network import Network +from .replayBuffer import PrioritizedReplayBuffer, ReplayBuffer + + +class DQNAgent: + """Categorical Double DQN with dueling NoisyNet, PER, and n-step replay. + + The high-level, simulator-specific loop lives in ``speed_training.py``. + This class owns action selection, transition collection, and optimizer + updates so it can also be reused by short integration checks. + """ + + def __init__( + self, + env, + memory_size: int, + batch_size: int, + target_update: int, + seed: int, + lr: float = 1e-4, + gamma: float = 0.97, + tau: float = 0.5, + frame_skip: int = 10, + epsilon: float = 1.0, + epsilon_decay: float = 0.999, + min_epsilon: float = 0.1, + hard_exploration_steps: int = 0, + exploration_steps: int = 0, + alpha: float = 0.2, + beta: float = 0.6, + prior_eps: float = 1e-6, + v_min: float = 0.0, + v_max: float = 120.0, + atom_size: int = 121, + n_step: int = 3, + n_step_alpha: float = 1.0, + hidden_dim: int = 256, + device=None, + **legacy_options, + ): + # Logging/checkpoint options belonged to the unrecovered monolithic + # trainer. Accept them for source compatibility; public loops own I/O. + ignored = {"log_dir", "file_path", "name", "ckpt_save_freq"} + unknown = set(legacy_options).difference(ignored) + if unknown: + raise TypeError(f"Unexpected DQN options: {', '.join(sorted(unknown))}") + if batch_size <= 0 or memory_size < batch_size: + raise ValueError("memory_size must be at least a positive batch_size") + if frame_skip <= 0 or atom_size < 2 or v_max <= v_min: + raise ValueError("Invalid frame skip or categorical support") + + random.seed(seed) + np.random.seed(seed) + self.env = env + self.frame_skip = int(frame_skip) + self.batch_size = int(batch_size) + self.target_update = int(target_update) + self.seed = int(seed) + self.gamma = float(gamma) + self.tau = float(tau) + self.device = torch.device( + device or ("cuda" if torch.cuda.is_available() else "cpu") + ) + + self.epsilon = float(epsilon) + self.epsilon_decay = float(epsilon_decay) + self.exploration_steps = int(exploration_steps) + self.hard_exploration_steps = int(hard_exploration_steps) + self.min_epsilon = float(min_epsilon) + self.beta = float(beta) + self.prior_eps = float(prior_eps) + + obs_dim = int(env.obs_space) + action_dim = int(env.action_space) + self.memory = PrioritizedReplayBuffer( + obs_dim, memory_size, batch_size, alpha=alpha, gamma=gamma + ) + self.use_n_step = n_step > 1 + self.n_step = int(n_step) + self.n_step_alpha = float(n_step_alpha) + if self.use_n_step: + self.memory_n = ReplayBuffer( + obs_dim, + memory_size, + batch_size, + n_step=n_step, + gamma=gamma, + ) + + self.v_min = float(v_min) + self.v_max = float(v_max) + self.atom_size = int(atom_size) + self.support = torch.linspace(v_min, v_max, atom_size).to(self.device) + self.dqn = Network( + obs_dim, action_dim, atom_size, self.support, hidden_dim=hidden_dim + ).to(self.device) + self.dqn_target = Network( + obs_dim, action_dim, atom_size, self.support, hidden_dim=hidden_dim + ).to(self.device) + self.dqn_target.load_state_dict(self.dqn.state_dict()) + self.dqn_target.eval() + self.optimizer = optim.Adam(self.dqn.parameters(), lr=lr) + self.transition = [] + self.is_test = False + + def select_action(self, state: np.ndarray) -> int: + """Select a discrete speed action and start a replay transition.""" + + explore = not self.is_test and np.random.uniform() < self.epsilon + if explore: + action = int(np.random.randint(self.env.action_space)) + else: + state_tensor = torch.as_tensor( + state, dtype=torch.float32, device=self.device + ).unsqueeze(0) + with torch.no_grad(): + action = int(self.dqn(state_tensor).argmax(dim=1).item()) + if not self.is_test: + self.transition = [np.asarray(state, dtype=np.float32).copy(), action] + return action + + def step(self, action: int, frame_skip: int | None = None): + """Execute one decision-level speed action and store one transition.""" + + repeats = self.frame_skip if frame_skip is None else int(frame_skip) + if repeats <= 0: + raise ValueError("frame_skip must be positive") + started = perf_counter() + step_decision = getattr(self.env, "step_decision", None) + if step_decision is not None: + next_state, total_reward, done, info = step_decision( + action, frame_skip=repeats + ) + else: + total_reward = 0.0 + info = {"success": False} + for _ in range(repeats): + next_state, reward, done, info = self.env.step(action) + total_reward += float(reward) + if done: + break + env_step_time = perf_counter() - started + + buffer_started = perf_counter() + if not self.is_test: + if len(self.transition) != 2: + raise RuntimeError("Call select_action() before step() while training") + transition = self.transition + [total_reward, next_state, done] + if self.use_n_step: + one_step_transition = self.memory_n.store(*transition) + else: + one_step_transition = transition + if one_step_transition: + self.memory.store(*one_step_transition) + buffer_time = perf_counter() - buffer_started + result_info = dict(info) + result_info.update( + env_step_time=env_step_time, + add_to_buffer_time=buffer_time, + ) + return next_state, total_reward, done, result_info + + def update_model(self) -> float: + """Perform one prioritized categorical DQN update.""" + + samples = self.memory.sample_batch(self.beta) + weights = torch.as_tensor( + samples["weights"].reshape(-1, 1), + dtype=torch.float32, + device=self.device, + ) + indices = samples["indices"] + elementwise_loss = self._compute_dqn_loss(samples, self.gamma) + if self.use_n_step: + n_step_samples = self.memory_n.sample_batch_from_idxs(indices) + elementwise_loss = elementwise_loss + self.n_step_alpha * self._compute_dqn_loss( + n_step_samples, self.gamma**self.n_step + ) + loss = torch.mean(elementwise_loss * weights) + + self.optimizer.zero_grad() + loss.backward() + clip_grad_norm_(self.dqn.parameters(), 10.0) + self.optimizer.step() + priorities = elementwise_loss.detach().cpu().numpy() + self.prior_eps + self.memory.update_priorities(indices, priorities) + self.dqn.reset_noise() + self.dqn_target.reset_noise() + return float(loss.item()) + + def decay_epsilon(self, step: int): + if step <= self.hard_exploration_steps: + self.epsilon = 1.0 + elif self.exploration_steps <= 0: + self.epsilon = 0.0 + elif step >= self.exploration_steps: + self.epsilon = self.min_epsilon + else: + self.epsilon = max(self.epsilon * self.epsilon_decay, self.min_epsilon) + + def set_eval(self, enabled=True): + self.is_test = bool(enabled) + self.dqn.eval() if enabled else self.dqn.train() + + def _compute_dqn_loss( + self, samples: Dict[str, np.ndarray], gamma: float + ) -> torch.Tensor: + state = torch.as_tensor(samples["obs"], dtype=torch.float32, device=self.device) + next_state = torch.as_tensor( + samples["next_obs"], dtype=torch.float32, device=self.device + ) + action = torch.as_tensor(samples["acts"], dtype=torch.long, device=self.device) + reward = torch.as_tensor( + samples["rews"].reshape(-1, 1), dtype=torch.float32, device=self.device + ) + done = torch.as_tensor( + samples["done"].reshape(-1, 1), dtype=torch.float32, device=self.device + ) + delta_z = (self.v_max - self.v_min) / (self.atom_size - 1) + + with torch.no_grad(): + next_action = self.dqn(next_state).argmax(1) + next_dist = self.dqn_target.dist(next_state)[ + range(self.batch_size), next_action + ] + target_support = ( + reward + (1 - done) * gamma * self.support + ).clamp(self.v_min, self.v_max) + projection = (target_support - self.v_min) / delta_z + lower = projection.floor().long() + upper = projection.ceil().long() + lower[(upper == lower) & (upper > 0)] -= 1 + upper[(upper == lower) & (lower < self.atom_size - 1)] += 1 + offset = ( + torch.arange(self.batch_size, device=self.device).unsqueeze(1) + * self.atom_size + ) + projected_dist = torch.zeros_like(next_dist) + projected_dist.view(-1).index_add_( + 0, + (lower + offset).view(-1), + (next_dist * (upper.float() - projection)).view(-1), + ) + projected_dist.view(-1).index_add_( + 0, + (upper + offset).view(-1), + (next_dist * (projection - lower.float())).view(-1), + ) + + distribution = self.dqn.dist(state) + log_probability = torch.log( + distribution[range(self.batch_size), action].clamp_min(1e-8) + ) + return -(projected_dist * log_probability).sum(1) + + def _target_hard_update(self): + self.dqn_target.load_state_dict(self.dqn.state_dict()) + + def _target_soft_update(self): + with torch.no_grad(): + for target, source in zip( + self.dqn_target.parameters(), self.dqn.parameters() + ): + target.copy_(self.tau * source + (1.0 - self.tau) * target) + + def save(self, path): + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + torch.save(self.dqn.state_dict(), path) + + def load(self, path): + self.dqn.load_state_dict( + torch.load(path, map_location=self.device, weights_only=True) + ) + self._target_hard_update() diff --git a/rl/rainbowDQN/network.py b/rl/rainbowDQN/network.py new file mode 100644 index 0000000..29fa77c --- /dev/null +++ b/rl/rainbowDQN/network.py @@ -0,0 +1,196 @@ +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +class NoisyLinear(nn.Module): + """Noisy linear module for NoisyNet. + + Attributes: + in_features (int): input size of linear module + out_features (int): output size of linear module + std_init (float): initial std value + weight_mu (nn.Parameter): mean value weight parameter + weight_sigma (nn.Parameter): std value weight parameter + bias_mu (nn.Parameter): mean value bias parameter + bias_sigma (nn.Parameter): std value bias parameter + + """ + + def __init__( + self, + in_features: int, + out_features: int, + std_init: float = 0.5, + ): + """Initialization.""" + super(NoisyLinear, self).__init__() + + self.in_features = in_features + self.out_features = out_features + self.std_init = std_init + + self.weight_mu = nn.Parameter(torch.Tensor(out_features, in_features)) + self.weight_sigma = nn.Parameter( + torch.Tensor(out_features, in_features) + ) + self.register_buffer( + "weight_epsilon", torch.Tensor(out_features, in_features) + ) + + self.bias_mu = nn.Parameter(torch.Tensor(out_features)) + self.bias_sigma = nn.Parameter(torch.Tensor(out_features)) + self.register_buffer("bias_epsilon", torch.Tensor(out_features)) + + self.reset_parameters() + self.reset_noise() + + def reset_parameters(self): + """Reset trainable network parameters (factorized gaussian noise).""" + mu_range = 1 / math.sqrt(self.in_features) + self.weight_mu.data.uniform_(-mu_range, mu_range) + self.weight_sigma.data.fill_( + self.std_init / math.sqrt(self.in_features) + ) + self.bias_mu.data.uniform_(-mu_range, mu_range) + self.bias_sigma.data.fill_( + self.std_init / math.sqrt(self.out_features) + ) + + def reset_noise(self): + """Make new noise.""" + epsilon_in = self.scale_noise(self.in_features) + epsilon_out = self.scale_noise(self.out_features) + + # outer product + self.weight_epsilon.copy_(epsilon_out.ger(epsilon_in)) + self.bias_epsilon.copy_(epsilon_out) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward method implementation. + + We don't use separate statements on train / eval mode. + It doesn't show remarkable difference of performance. + """ + if self.training: + weight = self.weight_mu + self.weight_sigma * self.weight_epsilon + bias = self.bias_mu + self.bias_sigma * self.bias_epsilon + else: + weight = self.weight_mu + bias = self.bias_mu + return F.linear(x, weight, bias) + + @staticmethod + def scale_noise(size: int) -> torch.Tensor: + """Set scale to make noise (factorized gaussian noise).""" + x = torch.randn(size) + + return x.sign().mul(x.abs().sqrt()) + + +class ImageBackbone(nn.Module): + def __init__(self, out_dim): + super(ImageBackbone, self).__init__() + import torchvision.models as models + #self.resnet18 = models.resnet18(pretrained=True) + #self.resnet18.fc = nn.Linear(512, out_dim) # Modify the fully connected layer for your desired number of classes + + #self.backbone = models.mobilenet_v2(pretrained=True) + #in_features = self.backbone.classifier[-1].in_features # Get the number of input features to the last layer + #self.backbone.classifier[-1] = nn.Linear(in_features, out_dim) + + self.backbone = models.squeezenet1_0(pretrained=True) + self.backbone.classifier[1] = nn.Conv2d(512, out_dim, kernel_size=(1, 1), stride=(1, 1)) # Adjust to your 'out_dim' + + def forward(self, x): + x = self.backbone(x) + return x + + +class Network(nn.Module): + def __init__( + self, + in_dim: int, + out_dim: int, + atom_size: int, + support: torch.Tensor, + hidden_dim: int = 256, + use_state = True + ): + """Initialization.""" + super(Network, self).__init__() + + self.support = support + self.in_dim = in_dim + self.out_dim = out_dim + self.hidden_dim = hidden_dim + self.atom_size = atom_size + self.use_state = use_state + + # Visual observations are encoded into finite feature vectors before + # reaching Rainbow, so the same MLP handles state-only and multimodal + # observation stacks. ``use_state`` remains for checkpoint compatibility. + self.feature_layer = nn.Sequential( + nn.Linear(in_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), + ) + + # set advantage layer + self.advantage_hidden_layer = NoisyLinear(hidden_dim, hidden_dim) + self.advantage_layer = NoisyLinear(hidden_dim, out_dim * atom_size) + + # set value layer + self.value_hidden_layer = NoisyLinear(hidden_dim, hidden_dim) + self.value_layer = NoisyLinear(hidden_dim, atom_size) + + # norm stats as non-learnable constants + self.states_mean = torch.nn.Parameter(torch.zeros(in_dim), requires_grad=False) + self.states_std = torch.nn.Parameter(torch.ones(in_dim), requires_grad=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward method implementation.""" + dist = self.dist(x) + q = torch.sum(dist * self.support, dim=2) + return q + + def update_norm_stats(self, norm_stats: dict): + with torch.no_grad(): + mean = torch.as_tensor( + norm_stats['states_mean'], dtype=self.states_mean.dtype, + device=self.states_mean.device, + ) + std = torch.as_tensor( + norm_stats['states_std'], dtype=self.states_std.dtype, + device=self.states_std.device, + ).clamp_min(1e-6) + self.states_mean.copy_(mean) + self.states_std.copy_(std) + + def dist(self, x: torch.Tensor) -> torch.Tensor: + """Get distribution for atoms.""" + x = (x - self.states_mean) / self.states_std.clamp_min(1e-6) + + feature = self.feature_layer(x) + + adv_hid = F.relu(self.advantage_hidden_layer(feature)) + val_hid = F.relu(self.value_hidden_layer(feature)) + + advantage = self.advantage_layer(adv_hid).view( + -1, self.out_dim, self.atom_size + ) + value = self.value_layer(val_hid).view(-1, 1, self.atom_size) + q_atoms = value + advantage - advantage.mean(dim=1, keepdim=True) + + dist = F.softmax(q_atoms, dim=-1).clamp_min(1e-6) + return dist / dist.sum(dim=-1, keepdim=True) + + def reset_noise(self): + """Reset all noisy layers.""" + self.advantage_hidden_layer.reset_noise() + self.advantage_layer.reset_noise() + self.value_hidden_layer.reset_noise() + self.value_layer.reset_noise() diff --git a/rl/rainbowDQN/replayBuffer.py b/rl/rainbowDQN/replayBuffer.py new file mode 100644 index 0000000..de38d0b --- /dev/null +++ b/rl/rainbowDQN/replayBuffer.py @@ -0,0 +1,234 @@ +import numpy as np +from typing import Deque, Dict, Tuple, List +from collections import deque +from .segment_tree import MinSegmentTree, SumSegmentTree +import random + +class ReplayBuffer: + """A simple numpy replay buffer.""" + + def __init__( + self, + obs_dim: int, + size: int, + batch_size: int = 32, + n_step: int = 1, + gamma: float = 0.99 + ): + self.obs_buf = np.zeros([size, obs_dim], dtype=np.float32) + self.next_obs_buf = np.zeros([size, obs_dim], dtype=np.float32) + self.acts_buf = np.zeros([size], dtype=np.float32) + self.rews_buf = np.zeros([size], dtype=np.float32) + self.done_buf = np.zeros(size, dtype=np.float32) + self.max_size, self.batch_size = size, batch_size + self.ptr, self.size, = 0, 0 + + # for N-step Learning + self.n_step_buffer = deque(maxlen=n_step) + self.n_step = n_step + self.gamma = gamma + + def store( + self, + obs: np.ndarray, + act: np.ndarray, + rew: float, + next_obs: np.ndarray, + done: bool, + ) -> Tuple[np.ndarray, np.ndarray, float, np.ndarray, bool]: + transition = (obs, act, rew, next_obs, done) + self.n_step_buffer.append(transition) + + # single step transition is not ready + if len(self.n_step_buffer) < self.n_step: + return () + + # make a n-step transition + rew, next_obs, done = self._get_n_step_info( + self.n_step_buffer, self.gamma + ) + obs, act = self.n_step_buffer[0][:2] + + self.obs_buf[self.ptr] = obs + self.next_obs_buf[self.ptr] = next_obs + self.acts_buf[self.ptr] = act + self.rews_buf[self.ptr] = rew + self.done_buf[self.ptr] = done + self.ptr = (self.ptr + 1) % self.max_size + self.size = min(self.size + 1, self.max_size) + + first_transition = self.n_step_buffer[0] + if done: + # Do not allow the retained n-step window to cross an episode reset. + # The final incomplete n-step tails are intentionally discarded so + # the one-step and n-step buffers retain matching indices. + self.n_step_buffer.clear() + return first_transition + + def sample_batch(self) -> Dict[str, np.ndarray]: + idxs = np.random.choice(self.size, size=self.batch_size, replace=False) + + return dict( + obs=self.obs_buf[idxs], + next_obs=self.next_obs_buf[idxs], + acts=self.acts_buf[idxs], + rews=self.rews_buf[idxs], + done=self.done_buf[idxs], + # for N-step Learning + indices=idxs, + ) + + def sample_batch_from_idxs( + self, idxs: np.ndarray + ) -> Dict[str, np.ndarray]: + # for N-step Learning + return dict( + obs=self.obs_buf[idxs], + next_obs=self.next_obs_buf[idxs], + acts=self.acts_buf[idxs], + rews=self.rews_buf[idxs], + done=self.done_buf[idxs], + ) + + def _get_n_step_info( + self, n_step_buffer: Deque, gamma: float + ) -> Tuple[np.int64, np.ndarray, bool]: + """Return n step rew, next_obs, and done.""" + # info of the last transition + rew, next_obs, done = n_step_buffer[-1][-3:] + + for transition in reversed(list(n_step_buffer)[:-1]): + r, n_o, d = transition[-3:] + + rew = r + gamma * rew * (1 - d) + next_obs, done = (n_o, d) if d else (next_obs, done) + + return rew, next_obs, done + + def __len__(self) -> int: + return self.size + + +class PrioritizedReplayBuffer(ReplayBuffer): + """Prioritized Replay buffer. + + Attributes: + max_priority (float): max priority + tree_ptr (int): next index of tree + alpha (float): alpha parameter for prioritized replay buffer + sum_tree (SumSegmentTree): sum tree for prior + min_tree (MinSegmentTree): min tree for min prior to get max weight + + """ + + def __init__( + self, + obs_dim: int, + size: int, + batch_size: int = 32, + alpha: float = 0.6, + n_step: int = 1, + gamma: float = 0.99, + ): + """Initialization.""" + assert alpha >= 0 + + super(PrioritizedReplayBuffer, self).__init__( + obs_dim, size, batch_size, n_step, gamma + ) + self.max_priority, self.tree_ptr = 1.0, 0 + self.alpha = alpha + + # capacity must be positive and a power of 2. + tree_capacity = 1 + while tree_capacity < self.max_size: + tree_capacity *= 2 + + self.sum_tree = SumSegmentTree(tree_capacity) + self.min_tree = MinSegmentTree(tree_capacity) + + def store( + self, + obs: np.ndarray, + act: int, + rew: float, + next_obs: np.ndarray, + done: bool, + ) -> Tuple[np.ndarray, np.ndarray, float, np.ndarray, bool]: + """Store experience and priority.""" + transition = super().store(obs, act, rew, next_obs, done) + + if transition: + self.sum_tree[self.tree_ptr] = self.max_priority ** self.alpha + self.min_tree[self.tree_ptr] = self.max_priority ** self.alpha + self.tree_ptr = (self.tree_ptr + 1) % self.max_size + + return transition + + def sample_batch(self, beta: float = 0.4) -> Dict[str, np.ndarray]: + """Sample a batch of experiences.""" + assert len(self) >= self.batch_size + assert beta > 0 + + indices = self._sample_proportional() + + obs = self.obs_buf[indices] + next_obs = self.next_obs_buf[indices] + acts = self.acts_buf[indices] + rews = self.rews_buf[indices] + done = self.done_buf[indices] + weights = np.array([self._calculate_weight(i, beta) for i in indices]) + + return dict( + obs=obs, + next_obs=next_obs, + acts=acts, + rews=rews, + done=done, + weights=weights, + indices=indices, + ) + + def update_priorities(self, indices: List[int], priorities: np.ndarray): + """Update priorities of sampled transitions.""" + assert len(indices) == len(priorities) + + for idx, priority in zip(indices, priorities): + assert priority > 0 + assert 0 <= idx < len(self) + + self.sum_tree[idx] = priority ** self.alpha + self.min_tree[idx] = priority ** self.alpha + + self.max_priority = max(self.max_priority, priority) + + def _sample_proportional(self) -> List[int]: + """Sample indices based on proportions.""" + indices = [] + # SegmentTree uses a half-open [start, end) range. Include the newest + # valid replay entry; the historical ``len(self) - 1`` bound silently + # excluded it from sampling. + p_total = self.sum_tree.sum(0, len(self)) + segment = p_total / self.batch_size + + for i in range(self.batch_size): + a = segment * i + b = segment * (i + 1) + upperbound = random.uniform(a, b) + idx = self.sum_tree.retrieve(upperbound) + indices.append(idx) + + return indices + + def _calculate_weight(self, idx: int, beta: float): + """Calculate the weight of the experience at idx.""" + # get max weight + p_min = self.min_tree.min() / self.sum_tree.sum() + max_weight = (p_min * len(self)) ** (-beta) + + # calculate weights + p_sample = self.sum_tree[idx] / self.sum_tree.sum() + weight = (p_sample * len(self)) ** (-beta) + weight = weight / max_weight + + return weight diff --git a/rl/rainbowDQN/segment_tree.py b/rl/rainbowDQN/segment_tree.py new file mode 100644 index 0000000..a19e49c --- /dev/null +++ b/rl/rainbowDQN/segment_tree.py @@ -0,0 +1,142 @@ +# -*- coding: utf-8 -*- +"""Segment tree for Prioritized Replay Buffer.""" + +import operator +from typing import Callable + + +class SegmentTree: + """ Create SegmentTree. + + Taken from OpenAI baselines github repository: + https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py + + Attributes: + capacity (int) + tree (list) + operation (function) + + """ + + def __init__(self, capacity: int, operation: Callable, init_value: float): + """Initialization. + + Args: + capacity (int) + operation (function) + init_value (float) + + """ + assert ( + capacity > 0 and capacity & (capacity - 1) == 0 + ), "capacity must be positive and a power of 2." + self.capacity = capacity + self.tree = [init_value for _ in range(2 * capacity)] + self.operation = operation + + def _operate_helper( + self, start: int, end: int, node: int, node_start: int, node_end: int + ) -> float: + """Returns result of operation in segment.""" + if start == node_start and end == node_end: + return self.tree[node] + mid = (node_start + node_end) // 2 + if end <= mid: + return self._operate_helper(start, end, 2 * node, node_start, mid) + else: + if mid + 1 <= start: + return self._operate_helper(start, end, 2 * node + 1, mid + 1, node_end) + else: + return self.operation( + self._operate_helper(start, mid, 2 * node, node_start, mid), + self._operate_helper(mid + 1, end, 2 * node + 1, mid + 1, node_end), + ) + + def operate(self, start: int = 0, end: int = 0) -> float: + """Returns result of applying `self.operation`.""" + if end <= 0: + end += self.capacity + end -= 1 + + return self._operate_helper(start, end, 1, 0, self.capacity - 1) + + def __setitem__(self, idx: int, val: float): + """Set value in tree.""" + idx += self.capacity + self.tree[idx] = val + + idx //= 2 + while idx >= 1: + self.tree[idx] = self.operation(self.tree[2 * idx], self.tree[2 * idx + 1]) + idx //= 2 + + def __getitem__(self, idx: int) -> float: + """Get real value in leaf node of tree.""" + assert 0 <= idx < self.capacity + + return self.tree[self.capacity + idx] + + +class SumSegmentTree(SegmentTree): + """ Create SumSegmentTree. + + Taken from OpenAI baselines github repository: + https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py + + """ + + def __init__(self, capacity: int): + """Initialization. + + Args: + capacity (int) + + """ + super(SumSegmentTree, self).__init__( + capacity=capacity, operation=operator.add, init_value=0.0 + ) + + def sum(self, start: int = 0, end: int = 0) -> float: + """Returns arr[start] + ... + arr[end].""" + return super(SumSegmentTree, self).operate(start, end) + + def retrieve(self, upperbound: float) -> int: + """Find the highest index `i` about upper bound in the tree""" + # TODO: Check assert case and fix bug + assert 0 <= upperbound <= self.sum() + 1e-5, "upperbound: {}".format(upperbound) + + idx = 1 + + while idx < self.capacity: # while non-leaf + left = 2 * idx + right = left + 1 + if self.tree[left] > upperbound: + idx = 2 * idx + else: + upperbound -= self.tree[left] + idx = right + return idx - self.capacity + + +class MinSegmentTree(SegmentTree): + """ Create SegmentTree. + + Taken from OpenAI baselines github repository: + https://github.com/openai/baselines/blob/master/baselines/common/segment_tree.py + + """ + + def __init__(self, capacity: int): + """Initialization. + + Args: + capacity (int) + + """ + super(MinSegmentTree, self).__init__( + capacity=capacity, operation=min, init_value=float("inf") + ) + + def min(self, start: int = 0, end: int = 0) -> float: + """Returns min(arr[start], ..., arr[end]).""" + return super(MinSegmentTree, self).operate(start, end) \ No newline at end of file diff --git a/scripted_policy.py b/scripted_policy.py new file mode 100644 index 0000000..57ab471 --- /dev/null +++ b/scripted_policy.py @@ -0,0 +1,233 @@ +import numpy as np +from pyquaternion import Quaternion + +from sim_tasks import normalize_task_name + + +class BasePolicy: + def __init__(self, inject_noise=False): + self.inject_noise = inject_noise + self.step_count = 0 + self.left_trajectory = None + self.right_trajectory = None + + def reset(self): + self.step_count = 0 + self.left_trajectory = None + self.right_trajectory = None + + def generate_trajectory(self, ts_first): + raise NotImplementedError + + @staticmethod + def interpolate(curr_waypoint, next_waypoint, t): + t_frac = (t - curr_waypoint["t"]) / (next_waypoint["t"] - curr_waypoint["t"] + 1e-8) + curr_xyz = curr_waypoint['xyz'] + curr_quat = curr_waypoint['quat'] + curr_grip = curr_waypoint['gripper'] + next_xyz = next_waypoint['xyz'] + next_quat = next_waypoint['quat'] + next_grip = next_waypoint['gripper'] + xyz = curr_xyz + (next_xyz - curr_xyz) * t_frac + quat = curr_quat + (next_quat - curr_quat) * t_frac + gripper = curr_grip + (next_grip - curr_grip) * t_frac + return xyz, quat, gripper + + def __call__(self, ts, step_inc=1): + if step_inc <= 0: + raise ValueError("step_inc must be positive") + + # generate trajectory at first timestep, then open-loop execution + if self.step_count == 0: + self.generate_trajectory(ts) + + curr_left_waypoint, next_left_waypoint = self._waypoint_pair( + self.left_trajectory, self.step_count + ) + curr_right_waypoint, next_right_waypoint = self._waypoint_pair( + self.right_trajectory, self.step_count + ) + + # interpolate between waypoints to obtain current pose and gripper command + left_xyz, left_quat, left_gripper = self.interpolate( + curr_left_waypoint, next_left_waypoint, self.step_count + ) + right_xyz, right_quat, right_gripper = self.interpolate( + curr_right_waypoint, next_right_waypoint, self.step_count + ) + + # Inject noise + if self.inject_noise: + scale = 0.01 + left_xyz = left_xyz + np.random.uniform(-scale, scale, left_xyz.shape) + right_xyz = right_xyz + np.random.uniform(-scale, scale, right_xyz.shape) + + action_left = np.concatenate([left_xyz, left_quat, [left_gripper]]) + action_right = np.concatenate([right_xyz, right_quat, [right_gripper]]) + + self.step_count += step_inc + return np.concatenate([action_left, action_right]) + + @staticmethod + def _waypoint_pair(trajectory, timestep): + """Find a safe interpolation bracket, including at the final waypoint.""" + + if not trajectory or len(trajectory) < 2: + raise ValueError("A scripted trajectory must contain at least two waypoints") + if timestep >= trajectory[-1]["t"]: + return trajectory[-1], trajectory[-1] + for current, following in zip(trajectory[:-1], trajectory[1:]): + if current["t"] <= timestep < following["t"]: + return current, following + return trajectory[0], trajectory[1] + + +class PickAndTransferPolicy(BasePolicy): + + def generate_trajectory(self, ts_first): + init_mocap_pose_right = ts_first.observation['mocap_pose_right'] + init_mocap_pose_left = ts_first.observation['mocap_pose_left'] + + box_info = np.array(ts_first.observation['env_state']) + box_xyz = box_info[:3] + box_quat = box_info[3:] + # print(f"Generate trajectory for {box_xyz=}") + + gripper_pick_quat = Quaternion(init_mocap_pose_right[3:]) + gripper_pick_quat = gripper_pick_quat * Quaternion(axis=[0.0, 1.0, 0.0], degrees=-60) + + meet_left_quat = Quaternion(axis=[1.0, 0.0, 0.0], degrees=90) + + meet_xyz = np.array([0, 0.5, 0.25]) + + self.left_trajectory = [ + {"t": 0, "xyz": init_mocap_pose_left[:3], "quat": init_mocap_pose_left[3:], "gripper": 0}, # sleep + {"t": 100, "xyz": meet_xyz + np.array([-0.1, 0, -0.02]), "quat": meet_left_quat.elements, "gripper": 1}, # approach meet position + {"t": 260, "xyz": meet_xyz + np.array([0.02, 0, -0.02]), "quat": meet_left_quat.elements, "gripper": 1}, # move to meet position + {"t": 310, "xyz": meet_xyz + np.array([0.02, 0, -0.02]), "quat": meet_left_quat.elements, "gripper": 0}, # close gripper + {"t": 360, "xyz": meet_xyz + np.array([-0.1, 0, -0.02]), "quat": np.array([1, 0, 0, 0]), "gripper": 0}, # move left + {"t": 400, "xyz": meet_xyz + np.array([-0.1, 0, -0.02]), "quat": np.array([1, 0, 0, 0]), "gripper": 0}, # stay + ] + + self.right_trajectory = [ + {"t": 0, "xyz": init_mocap_pose_right[:3], "quat": init_mocap_pose_right[3:], "gripper": 0}, # sleep + {"t": 90, "xyz": box_xyz + np.array([0, 0, 0.08]), "quat": gripper_pick_quat.elements, "gripper": 1}, # approach the cube + {"t": 130, "xyz": box_xyz + np.array([0., 0, -0.015]), "quat": gripper_pick_quat.elements, "gripper": 1}, # go down + {"t": 170, "xyz": box_xyz + np.array([0, 0, -0.015]), "quat": gripper_pick_quat.elements, "gripper": 0}, # close gripper + {"t": 200, "xyz": meet_xyz + np.array([0.05, 0, 0]), "quat": gripper_pick_quat.elements, "gripper": 0}, # approach meet position + {"t": 220, "xyz": meet_xyz, "quat": gripper_pick_quat.elements, "gripper": 0}, # move to meet position + {"t": 310, "xyz": meet_xyz, "quat": gripper_pick_quat.elements, "gripper": 1}, # open gripper + {"t": 360, "xyz": meet_xyz + np.array([0.1, 0, 0]), "quat": gripper_pick_quat.elements, "gripper": 1}, # move to right + {"t": 400, "xyz": meet_xyz + np.array([0.1, 0, 0]), "quat": gripper_pick_quat.elements, "gripper": 1}, # stay + ] + +class PickAndTransferTeaBagPolicy(BasePolicy): + + def generate_trajectory(self, ts_first): + init_mocap_pose_right = ts_first.observation['mocap_pose_right'] + init_mocap_pose_left = ts_first.observation['mocap_pose_left'] + + box_info = np.array(ts_first.observation['env_state']) + + box_xyz = box_info[:3] + box_quat = box_info[3:] + #print(f"Generate trajectory for {box_xyz=}") + + gripper_pick_quat_org = Quaternion(init_mocap_pose_right[3:]) + gripper_pick_quat = gripper_pick_quat_org * Quaternion(axis=[0.0, 1.0, 0.0], degrees=-90) + gripper_move_quat = gripper_pick_quat_org * Quaternion(axis=[0.0, 1.0, 0.0], degrees=-20) + + meet_left_quat = Quaternion(axis=[1.0, 0.0, 0.0], degrees=90) + + meet_xyz = np.array([-0.1, 0.6, 0.30]) + + + self.left_trajectory = [ + {"t": 0, "xyz": init_mocap_pose_left[:3], "quat": init_mocap_pose_left[3:], "gripper": 0}, # sleep + {"t": 500, "xyz": init_mocap_pose_left[:3], "quat": np.array([1, 0, 0, 0]), "gripper": 0}, # stay + ] + + ''' + # Old policy, oscillation is a problem + self.right_trajectory = [ + {"t": 0, "xyz": init_mocap_pose_right[:3], "quat": init_mocap_pose_right[3:], "gripper": 0}, # sleep + {"t": 50, "xyz": box_xyz + np.array([0, 0, 0.08]), "quat": gripper_pick_quat.elements, "gripper": 1}, # approach the cube + {"t": 70, "xyz": box_xyz + np.array([0.005, 0, -0.03]), "quat": gripper_pick_quat.elements, "gripper": 1}, # go down + {"t": 100, "xyz": box_xyz + np.array([0.005, 0, -0.03]), "quat": gripper_pick_quat.elements, "gripper": 0}, # close gripper + {"t": 170, "xyz": meet_xyz + np.array([0.05, 0, 0]), "quat": gripper_move_quat.elements, "gripper": 0}, # approach meet position + {"t": 200, "xyz": meet_xyz, "quat": gripper_move_quat.elements, "gripper": 0}, # move to meet position + {"t": 420, "xyz": meet_xyz, "quat": gripper_move_quat.elements, "gripper": 0}, # open gripper + {"t": 460, "xyz": meet_xyz, "quat": gripper_move_quat.elements, "gripper": 1}, # open gripper + {"t": 500, "xyz": meet_xyz + np.array([0.1, 0, 0]), "quat": gripper_pick_quat.elements, "gripper": 1}, + ] + ''' + + self.right_trajectory = [ + {"t": 0, "xyz": init_mocap_pose_right[:3], "quat": init_mocap_pose_right[3:], "gripper": 0}, # sleep + {"t": 50, "xyz": box_xyz + np.array([0, 0, 0.08]), "quat": gripper_pick_quat.elements, "gripper": 1}, # approach the cube + {"t": 70, "xyz": box_xyz + np.array([0.005, 0, -0.03]), "quat": gripper_pick_quat.elements, "gripper": 1}, # go down + {"t": 100, "xyz": box_xyz + np.array([0.005, 0, -0.03]), "quat": gripper_pick_quat.elements, "gripper": 0}, # close gripper + {"t": 150, "xyz": box_xyz + np.array([0.1, 0, 0.1]), "quat": gripper_pick_quat.elements, "gripper": 0}, # vertical align + {"t": 250, "xyz": box_xyz + np.array([0.1, 0, 0.3]), "quat": gripper_move_quat.elements, "gripper": 0}, # rise up + {"t": 400, "xyz": meet_xyz, "quat": gripper_move_quat.elements, "gripper": 0}, # move to meet position + {"t": 420, "xyz": meet_xyz, "quat": gripper_move_quat.elements, "gripper": 0}, # open gripper + {"t": 450, "xyz": meet_xyz, "quat": gripper_move_quat.elements, "gripper": 1}, # open gripper + {"t": 500, "xyz": init_mocap_pose_right[:3], "quat": gripper_move_quat.elements, "gripper": 1}, + ] + +class InsertionPolicy(BasePolicy): + + def generate_trajectory(self, ts_first): + init_mocap_pose_right = ts_first.observation['mocap_pose_right'] + init_mocap_pose_left = ts_first.observation['mocap_pose_left'] + + peg_info = np.array(ts_first.observation['env_state'])[:7] + peg_xyz = peg_info[:3] + peg_quat = peg_info[3:] + + socket_info = np.array(ts_first.observation['env_state'])[7:] + socket_xyz = socket_info[:3] + socket_quat = socket_info[3:] + + gripper_pick_quat_right = Quaternion(init_mocap_pose_right[3:]) + gripper_pick_quat_right = gripper_pick_quat_right * Quaternion(axis=[0.0, 1.0, 0.0], degrees=-60) + + gripper_pick_quat_left = Quaternion(init_mocap_pose_right[3:]) + gripper_pick_quat_left = gripper_pick_quat_left * Quaternion(axis=[0.0, 1.0, 0.0], degrees=60) + + meet_xyz = np.array([0, 0.5, 0.15]) + lift_right = 0.00715 + + self.left_trajectory = [ + {"t": 0, "xyz": init_mocap_pose_left[:3], "quat": init_mocap_pose_left[3:], "gripper": 0}, # sleep + {"t": 120, "xyz": socket_xyz + np.array([0, 0, 0.08]), "quat": gripper_pick_quat_left.elements, "gripper": 1}, # approach the cube + {"t": 170, "xyz": socket_xyz + np.array([0, 0, -0.03]), "quat": gripper_pick_quat_left.elements, "gripper": 1}, # go down + {"t": 220, "xyz": socket_xyz + np.array([0, 0, -0.03]), "quat": gripper_pick_quat_left.elements, "gripper": 0}, # close gripper + {"t": 285, "xyz": meet_xyz + np.array([-0.1, 0, 0]), "quat": gripper_pick_quat_left.elements, "gripper": 0}, # approach meet position + {"t": 340, "xyz": meet_xyz + np.array([-0.05, 0, 0]), "quat": gripper_pick_quat_left.elements,"gripper": 0}, # insertion + {"t": 400, "xyz": meet_xyz + np.array([-0.05, 0, 0]), "quat": gripper_pick_quat_left.elements, "gripper": 0}, # insertion + ] + + self.right_trajectory = [ + {"t": 0, "xyz": init_mocap_pose_right[:3], "quat": init_mocap_pose_right[3:], "gripper": 0}, # sleep + {"t": 120, "xyz": peg_xyz + np.array([0, 0, 0.08]), "quat": gripper_pick_quat_right.elements, "gripper": 1}, # approach the cube + {"t": 170, "xyz": peg_xyz + np.array([0, 0, -0.03]), "quat": gripper_pick_quat_right.elements, "gripper": 1}, # go down + {"t": 220, "xyz": peg_xyz + np.array([0, 0, -0.03]), "quat": gripper_pick_quat_right.elements, "gripper": 0}, # close gripper + {"t": 285, "xyz": meet_xyz + np.array([0.1, 0, lift_right]), "quat": gripper_pick_quat_right.elements, "gripper": 0}, # approach meet position + {"t": 340, "xyz": meet_xyz + np.array([0.05, 0, lift_right]), "quat": gripper_pick_quat_right.elements, "gripper": 0}, # insertion + {"t": 400, "xyz": meet_xyz + np.array([0.05, 0, lift_right]), "quat": gripper_pick_quat_right.elements, "gripper": 0}, # insertion + + ] + + +POLICY_CLASSES = { + "pick_and_place": PickAndTransferPolicy, + "insertion": InsertionPolicy, + "tea_bag": PickAndTransferTeaBagPolicy, +} + + +def make_scripted_policy(task_name, inject_noise=False): + """Construct the matching historical scripted policy for a sim task.""" + + return POLICY_CLASSES[normalize_task_name(task_name)](inject_noise=inject_noise) diff --git a/scripts/__init__.py b/scripts/__init__.py new file mode 100644 index 0000000..2e7511c --- /dev/null +++ b/scripts/__init__.py @@ -0,0 +1 @@ +"""Command-line entry points for the SpeedTuning simulator release.""" diff --git a/scripts/check_chunked_policy.py b/scripts/check_chunked_policy.py new file mode 100644 index 0000000..d70a822 --- /dev/null +++ b/scripts/check_chunked_policy.py @@ -0,0 +1,25 @@ +#!/usr/bin/env python3 +"""Exercise ACT-style action chunks against every joint simulator.""" + +import json +import sys +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from chunked_policy import replay_recorded_chunks # noqa: E402 +from sim_tasks import TASK_SPECS # noqa: E402 + + +def main(): + results = [replay_recorded_chunks(task) for task in TASK_SPECS] + for result in results: + print(json.dumps(result, sort_keys=True)) + return 0 if all(result["success"] for result in results) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/eval_speed_policy.py b/scripts/eval_speed_policy.py new file mode 100644 index 0000000..34fbc13 --- /dev/null +++ b/scripts/eval_speed_policy.py @@ -0,0 +1,147 @@ +#!/usr/bin/env python3 +"""Evaluate fixed, profiled, Rainbow, or external speed policies.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import numpy as np + + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from experiment_config import defaults_from_argv # noqa: E402 +from policy_speed_env import make_speed_reward # noqa: E402 +from policy_loader import load_speed_policy # noqa: E402 +from scripts.policy_cli import ( # noqa: E402 + add_base_policy_arguments, + add_observation_arguments, + build_speed_env, + comma_floats, + comma_ints, + json_object, +) +from speed_policy import ( # noqa: E402 + FixedSpeedPolicy, + RainbowSpeedPolicy, + SpeedProfilePolicy, + rollout_speed_policy, + summarize_rollouts, +) + + +def parse_args(): + config_defaults, config_metadata = defaults_from_argv() + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", help="JSON manifest or named preset such as paper-sim.") + parser.add_argument("--task", choices=("pick_and_place", "insertion", "tea_bag"), default="tea_bag") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--device", default="cpu") + parser.add_argument("--episodes", type=int, default=10) + parser.add_argument( + "--seeds", + type=comma_ints, + help="Explicit episode seeds; when set, one rollout is run per seed.", + ) + parser.add_argument("--frame-skip", type=int, default=10) + parser.add_argument("--success-bonus", type=float, default=100.0) + parser.add_argument("--speed-weight", type=float, default=0.01) + parser.add_argument("--speed-power", type=float, default=2.0) + parser.add_argument("--speed-values", type=comma_floats, default=(1.0, 1.5, 2.0, 2.5, 3.0)) + parser.add_argument( + "--speed-policy", + choices=("fixed", "profile", "rainbow", "external"), + default="fixed", + ) + parser.add_argument("--speed", type=float, default=1.0) + parser.add_argument("--profile", type=comma_floats) + parser.add_argument("--speed-checkpoint", type=Path) + parser.add_argument("--speed-policy-loader", help="External speed factory as module:attribute.") + parser.add_argument("--speed-factory-kwargs", type=json_object, default={}) + parser.add_argument("--video", type=Path) + add_base_policy_arguments(parser) + add_observation_arguments(parser) + parser.set_defaults(**config_defaults) + return parser.parse_args(), config_metadata + + +def build_speed_policy(args): + if args.speed_policy == "fixed": + return FixedSpeedPolicy(args.speed) + if args.speed_policy == "profile": + if args.profile is None: + raise ValueError("--profile is required for a profile speed policy") + return SpeedProfilePolicy(args.profile) + if args.speed_policy == "rainbow": + if args.speed_checkpoint is None: + raise ValueError("--speed-checkpoint is required for a Rainbow speed policy") + return RainbowSpeedPolicy.load(args.speed_checkpoint, device=args.device) + if not args.speed_policy_loader: + raise ValueError("--speed-policy-loader is required for an external speed policy") + return load_speed_policy( + args.speed_policy_loader, + checkpoint=args.speed_checkpoint, + device=args.device, + factory_kwargs=args.speed_factory_kwargs, + ) + + +def main(): + args, config_metadata = parse_args() + if args.episodes <= 0: + print("error: --episodes must be positive", file=sys.stderr) + return 2 + try: + policy = build_speed_policy(args) + reward_fn = make_speed_reward( + success_bonus=args.success_bonus, + speed_weight=args.speed_weight, + speed_power=args.speed_power, + ) + args.restore_observation_encoder = isinstance(policy, RainbowSpeedPolicy) + seeds = args.seeds or tuple(range(args.seed, args.seed + args.episodes)) + rollouts = [] + for index, seed in enumerate(seeds): + env = build_speed_env( + args, + reward_fn=reward_fn, + video_path=args.video if index == 0 else None, + seed=seed, + ) + try: + rollout = rollout_speed_policy( + env, + policy, + frame_skip=( + policy.frame_skip + if isinstance(policy, RainbowSpeedPolicy) + else args.frame_skip + ), + ) + rollout["seed"] = seed + rollouts.append(rollout) + finally: + env.close() + result = { + "task": args.task, + "base_policy": args.base_policy, + "speed_policy": args.speed_policy, + "seeds": list(seeds), + "experiment_config": config_metadata, + **summarize_rollouts(rollouts), + "rollouts": rollouts, + } + print(json.dumps(result, sort_keys=True)) + return 0 + except (ImportError, RuntimeError, ValueError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/policy_cli.py b/scripts/policy_cli.py new file mode 100644 index 0000000..f37243e --- /dev/null +++ b/scripts/policy_cli.py @@ -0,0 +1,179 @@ +"""Shared command-line helpers for base-policy and speed-policy tools.""" + +from __future__ import annotations + +import json + +from policy_loader import load_chunk_predictor, load_observation_encoder +from policy_speed_env import create_recorded_chunk_speed_env, create_speed_env +from speed_observation import StateObservationEncoder, VisualObservationEncoder + + +def comma_floats(value): + try: + values = tuple(float(item.strip()) for item in value.split(",")) + except ValueError as exc: + raise ValueError("Expected comma-separated numbers") from exc + if not values: + raise ValueError("Expected at least one number") + return values + + +def comma_strings(value): + values = tuple(item.strip() for item in value.split(",") if item.strip()) + if not values: + raise ValueError("Expected at least one comma-separated value") + return values + + +def comma_ints(value): + try: + values = tuple(int(item.strip()) for item in value.split(",")) + except ValueError as exc: + raise ValueError("Expected comma-separated integers") from exc + if not values: + raise ValueError("Expected at least one integer") + return values + + +def json_object(value): + parsed = json.loads(value) + if not isinstance(parsed, dict): + raise ValueError("Factory kwargs must be a JSON object") + return parsed + + +def add_base_policy_arguments(parser): + parser.add_argument( + "--base-policy", + choices=("scripted", "recorded-chunk", "external-chunk"), + default="scripted", + help="Robot policy wrapped by the speed controller.", + ) + parser.add_argument( + "--chunk-policy", + help="External chunk factory as module:attribute (required for external-chunk).", + ) + parser.add_argument("--upstream-checkpoint") + parser.add_argument("--factory-kwargs", type=json_object, default={}) + parser.add_argument("--chunk-size", type=int, default=25) + parser.add_argument( + "--no-policy-images", + action="store_true", + help="Do not render cameras for an external policy that only uses state.", + ) + parser.add_argument( + "--randomize-object-pose", + action="store_true", + help="Sample a new tea-bag pose on every reset (cube/insertion already vary).", + ) + + +def add_observation_arguments(parser): + parser.add_argument( + "--speed-observation", + choices=("state", "visual", "external"), + default="state", + help="Input representation used by the speed policy.", + ) + parser.add_argument("--frame-stack", type=int, default=1) + parser.add_argument( + "--camera-names", + type=comma_strings, + default=("top", "angle", "vis"), + ) + parser.add_argument( + "--image-encoder", + choices=("resnet18-pretrained", "resnet18-random"), + default="resnet18-pretrained", + ) + parser.add_argument("--image-size", type=int, default=224) + parser.add_argument("--observation-encoder-loader") + parser.add_argument("--observation-factory-kwargs", type=json_object, default={}) + parser.add_argument( + "--include-env-state", dest="include_env_state", action="store_true", default=True + ) + parser.add_argument("--no-env-state", dest="include_env_state", action="store_false") + parser.add_argument( + "--include-qpos", dest="include_qpos", action="store_true", default=True + ) + parser.add_argument("--no-qpos", dest="include_qpos", action="store_false") + parser.add_argument( + "--include-qvel", dest="include_qvel", action="store_true", default=True + ) + parser.add_argument("--no-qvel", dest="include_qvel", action="store_false") + + +def build_observation_encoder(args): + if args.speed_observation == "state": + return StateObservationEncoder( + include_qpos=args.include_qpos, + include_qvel=args.include_qvel, + include_env_state=args.include_env_state, + ) + if args.speed_observation == "visual": + return VisualObservationEncoder( + camera_names=args.camera_names, + pretrained=args.image_encoder == "resnet18-pretrained", + image_size=args.image_size, + device=args.device, + include_qpos=args.include_qpos, + include_qvel=args.include_qvel, + include_env_state=args.include_env_state, + initialize_pretrained=not getattr( + args, "restore_observation_encoder", False + ), + ) + if not args.observation_encoder_loader: + raise ValueError( + "--observation-encoder-loader is required for external observations" + ) + return load_observation_encoder( + args.observation_encoder_loader, + task_name=args.task, + checkpoint=args.upstream_checkpoint, + device=args.device, + factory_kwargs=args.observation_factory_kwargs, + ) + + +def build_speed_env(args, reward_fn=None, video_path=None, seed=None): + observation_encoder = build_observation_encoder(args) + common = { + "task_name": args.task, + "reward_fn": reward_fn, + "seed": args.seed if seed is None else seed, + "speed_values": args.speed_values, + "observation_encoder": observation_encoder, + "frame_stack": args.frame_stack, + "decision_frame_skip": args.frame_skip, + "randomize_object_pose": args.randomize_object_pose, + } + if video_path is not None: + common.update(save_video=True, video_path=video_path) + + if args.base_policy == "scripted": + return create_speed_env(**common) + if args.base_policy == "recorded-chunk": + return create_recorded_chunk_speed_env( + chunk_size=args.chunk_size, + render_images=getattr(observation_encoder, "requires_images", False), + **common, + ) + if not args.chunk_policy: + raise ValueError("--chunk-policy is required with --base-policy external-chunk") + predictor = load_chunk_predictor( + args.chunk_policy, + task_name=args.task, + checkpoint=args.upstream_checkpoint, + device=args.device, + factory_kwargs=args.factory_kwargs, + ) + return create_speed_env( + chunk_predictor=predictor, + render_images=( + not args.no_policy_images + or getattr(observation_encoder, "requires_images", False) + ), + **common, + ) diff --git a/scripts/rainbow_poc.py b/scripts/rainbow_poc.py new file mode 100644 index 0000000..5284bc8 --- /dev/null +++ b/scripts/rainbow_poc.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +"""Run a small Rainbow DQN optimization proof on the tea-bag speed task.""" + +from __future__ import annotations + +import argparse +import json +import random +import sys +from pathlib import Path + +import numpy as np + + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +try: + import torch +except ImportError as exc: + raise SystemExit("Rainbow POC requires: uv sync --extra rl") from exc + +from policy_speed_env import create_speed_env # noqa: E402 +from rl.rainbowDQN.dqnAgent import DQNAgent # noqa: E402 + + +def speed_reward(speed, done, success): + reward = speed**2 / 100.0 + if done and success: + reward += 100.0 + return reward + + +def run_poc(seed=0, min_transitions=64, updates=8): + """Collect real simulator transitions and perform Rainbow optimizer steps.""" + + np.random.seed(seed) + random.seed(seed) + torch.manual_seed(seed) + env = create_speed_env( + task_name="tea_bag", + reward_fn=speed_reward, + seed=seed, + ) + agent = DQNAgent( + env, + memory_size=2048, + batch_size=16, + target_update=4, + seed=seed, + lr=3e-4, + gamma=0.97, + frame_skip=5, + epsilon=1.0, + exploration_steps=1000, + min_epsilon=0.1, + atom_size=51, + v_min=0.0, + v_max=120.0, + n_step=3, + hidden_dim=64, + device="cpu", + log_dir=None, + ) + + episodes = 0 + successes = 0 + actions_seen = set() + rng = np.random.RandomState(seed) + decision_index = 0 + while len(agent.memory) < min_transitions: + state = env.reset() + done = False + info = {"success": False} + while not done: + # Exercise every action during the initial approach, then collect a + # predominantly safe 1.0x/1.5x rollout so terminal success is present. + if decision_index < env.action_space: + action = decision_index + else: + action = int(rng.choice(2, p=[0.8, 0.2])) + decision_index += 1 + actions_seen.add(action) + agent.transition = [state, action] + state, _, done, info = agent.step(action, frame_skip=agent.frame_skip) + episodes += 1 + successes += int(info["success"]) + + states = agent.memory.obs_buf[: len(agent.memory)] + norm_stats = { + "states_mean": states.mean(axis=0), + "states_std": states.std(axis=0), + } + agent.dqn.update_norm_stats(norm_stats) + agent.dqn_target.update_norm_stats(norm_stats) + + parameters_before = [parameter.detach().clone() for parameter in agent.dqn.parameters()] + losses = [] + for update in range(updates): + loss = agent.update_model() + losses.append(loss) + if (update + 1) % agent.target_update == 0: + agent._target_soft_update() + + parameter_delta = sum( + (before - after.detach()).abs().sum().item() + for before, after in zip(parameters_before, agent.dqn.parameters()) + ) + probe = torch.as_tensor(states[:4], dtype=torch.float32) + with torch.inference_mode(): + q_values = agent.dqn(probe) + distributions = agent.dqn.dist(probe) + + result = { + "task": "tea_bag", + "episodes": episodes, + "successful_episodes": successes, + "replay_transitions": len(agent.memory), + "actions_seen": sorted(actions_seen), + "updates": updates, + "loss_first": losses[0], + "loss_last": losses[-1], + "losses_finite": bool(np.isfinite(losses).all()), + "parameter_delta": parameter_delta, + "q_shape": list(q_values.shape), + "distributions_normalized": bool( + torch.allclose( + distributions.sum(dim=-1), + torch.ones_like(distributions.sum(dim=-1)), + atol=1e-5, + ) + ), + } + result["passed"] = bool( + result["losses_finite"] + and parameter_delta > 0 + and result["distributions_normalized"] + and len(agent.memory) >= min_transitions + ) + env.close() + return result + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--min-transitions", type=int, default=64) + parser.add_argument("--updates", type=int, default=8) + args = parser.parse_args() + result = run_poc(args.seed, args.min_transitions, args.updates) + print(json.dumps(result, sort_keys=True)) + return 0 if result["passed"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/run_sim.py b/scripts/run_sim.py new file mode 100644 index 0000000..3e2d395 --- /dev/null +++ b/scripts/run_sim.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Run the recovered SpeedTuning scripted simulation tasks.""" + +from __future__ import annotations + +import argparse +import json +import math +import sys +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from ee_sim_env import make_ee_sim_env # noqa: E402 +from scripted_policy import make_scripted_policy # noqa: E402 +from sim_tasks import TASK_SPECS, normalize_task_name # noqa: E402 + + +def run_task(task_name, speed=1.0, seed=0, video_path=None): + """Run one open-loop scripted rollout and return a JSON-safe summary.""" + + task_name = normalize_task_name(task_name) + spec = TASK_SPECS[task_name] + render_images = video_path is not None + env = make_ee_sim_env(task_name, render_images=render_images, seed=seed) + timestep = env.reset() + policy = make_scripted_policy(task_name) + rewards = [] + frames = [] + + num_steps = math.ceil(spec.episode_len / speed) + for _ in range(num_steps): + timestep = env.step(policy(timestep, step_inc=speed)) + rewards.append(int(timestep.reward or 0)) + if render_images: + frames.append(timestep.observation["images"]["angle"]) + + max_reward = max(rewards, default=0) + result = { + "task": task_name, + "success": max_reward == env.task.max_reward, + "max_reward": max_reward, + "target_reward": env.task.max_reward, + "return": sum(rewards), + "steps": num_steps, + "speed": speed, + "seed": seed, + } + + if video_path is not None: + try: + import imageio.v2 as imageio + except ImportError as exc: + raise RuntimeError( + "Video export requires: uv sync --extra video" + ) from exc + video_path.parent.mkdir(parents=True, exist_ok=True) + imageio.mimsave(video_path, frames, fps=max(1, round(50 / speed))) + result["video"] = str(video_path) + return result + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Run recovered pick-and-place, insertion, and tea-bag simulations." + ) + parser.add_argument( + "--task", + action="append", + help="Task to run; repeat the flag for multiple tasks (default: all).", + ) + parser.add_argument("--speed", type=float, default=1.0, help="Waypoint time increment.") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--video-dir", + type=Path, + help="Optional directory for angle-camera MP4 rollouts.", + ) + return parser.parse_args() + + +def main(): + args = parse_args() + if args.speed <= 0: + raise SystemExit("--speed must be positive") + tasks = args.task or list(TASK_SPECS) + results = [] + for requested_task in tasks: + task_name = normalize_task_name(requested_task) + video_path = None + if args.video_dir is not None: + video_path = args.video_dir / f"{task_name}.mp4" + result = run_task( + task_name, + speed=args.speed, + seed=args.seed, + video_path=video_path, + ) + results.append(result) + print(json.dumps(result, sort_keys=True)) + return 0 if all(result["success"] for result in results) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/sweep_speed_policy.py b/scripts/sweep_speed_policy.py new file mode 100644 index 0000000..42d9891 --- /dev/null +++ b/scripts/sweep_speed_policy.py @@ -0,0 +1,158 @@ +#!/usr/bin/env python3 +"""Run the paper-style fixed-speed sweep and optional adaptive-policy evaluation.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from experiment_config import defaults_from_argv # noqa: E402 +from policy_loader import load_speed_policy # noqa: E402 +from policy_speed_env import make_speed_reward # noqa: E402 +from scripts.policy_cli import ( # noqa: E402 + add_base_policy_arguments, + add_observation_arguments, + build_speed_env, + comma_floats, + comma_ints, + json_object, +) +from speed_evaluation import ( # noqa: E402 + evaluate_fixed_speed_sweep, + evaluate_seeded_policy, + plot_speed_success_tradeoff, + speed_grid, +) +from speed_policy import RainbowSpeedPolicy, SpeedProfilePolicy # noqa: E402 + + +def parse_args(): + defaults, config_metadata = defaults_from_argv() + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", help="JSON manifest or named preset such as paper-sim.") + parser.add_argument("--task", choices=("pick_and_place", "insertion", "tea_bag"), default="tea_bag") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--seeds", type=comma_ints) + parser.add_argument("--device", default="cpu") + parser.add_argument("--frame-skip", type=int, default=10) + parser.add_argument("--success-bonus", type=float, default=100.0) + parser.add_argument("--speed-weight", type=float, default=0.01) + parser.add_argument("--speed-power", type=float, default=2.0) + parser.add_argument("--speed-values", type=comma_floats, default=(1.0, 1.5, 2.0, 2.5, 3.0)) + parser.add_argument("--speed-start", type=float, default=1.0) + parser.add_argument("--speed-stop", type=float, default=4.5) + parser.add_argument("--speed-step", type=float, default=0.1) + parser.add_argument("--episodes-per-speed", type=int, default=100) + parser.add_argument("--adaptive-episodes", type=int, default=0) + parser.add_argument( + "--adaptive-policy", + choices=("rainbow", "external", "profile"), + ) + parser.add_argument("--profile", type=comma_floats) + parser.add_argument("--speed-checkpoint", type=Path) + parser.add_argument("--speed-policy-loader") + parser.add_argument("--speed-factory-kwargs", type=json_object, default={}) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--plot", type=Path) + add_base_policy_arguments(parser) + add_observation_arguments(parser) + parser.set_defaults(**defaults) + return parser.parse_args(), config_metadata + + +def build_adaptive_policy(args): + if args.adaptive_policy is None: + return None + if args.adaptive_policy == "rainbow": + if args.speed_checkpoint is None: + raise ValueError("--speed-checkpoint is required for Rainbow evaluation") + return RainbowSpeedPolicy.load(args.speed_checkpoint, device=args.device) + if args.adaptive_policy == "profile": + if args.profile is None: + raise ValueError("--profile is required for profile evaluation") + return SpeedProfilePolicy(args.profile) + if not args.speed_policy_loader: + raise ValueError("--speed-policy-loader is required for external evaluation") + return load_speed_policy( + args.speed_policy_loader, + checkpoint=args.speed_checkpoint, + device=args.device, + factory_kwargs=args.speed_factory_kwargs, + ) + + +def main(): + args, config_metadata = parse_args() + try: + if args.episodes_per_speed <= 0 or args.adaptive_episodes < 0: + raise ValueError("Episode counts must be positive (or zero for adaptive)") + adaptive_policy = build_adaptive_policy(args) + reward_fn = make_speed_reward( + success_bonus=args.success_bonus, + speed_weight=args.speed_weight, + speed_power=args.speed_power, + ) + args.restore_observation_encoder = isinstance( + adaptive_policy, RainbowSpeedPolicy + ) + fixed_seeds = args.seeds or tuple( + range(args.seed, args.seed + args.episodes_per_speed) + ) + + def env_factory(seed): + return build_speed_env(args, reward_fn=reward_fn, seed=seed) + + report = { + "task": args.task, + "base_policy": args.base_policy, + "frame_skip": args.frame_skip, + "metric": "physical_acceleration=episode_len/physics_steps", + "experiment_config": config_metadata, + "fixed_speed_sweep": evaluate_fixed_speed_sweep( + env_factory, + speed_grid(args.speed_start, args.speed_stop, args.speed_step), + fixed_seeds, + frame_skip=args.frame_skip, + ), + } + if adaptive_policy is not None: + count = args.adaptive_episodes or 2000 + adaptive_seeds = tuple(range(args.seed, args.seed + count)) + report["adaptive_policy"] = evaluate_seeded_policy( + env_factory, + adaptive_policy, + adaptive_seeds, + frame_skip=( + adaptive_policy.frame_skip + if isinstance(adaptive_policy, RainbowSpeedPolicy) + else args.frame_skip + ), + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") + if args.plot is not None: + plot_speed_success_tradeoff(report, args.plot) + summary = { + "output": str(args.output), + "plot": None if args.plot is None else str(args.plot), + "fixed_speed_points": len(report["fixed_speed_sweep"]), + "adaptive_episodes": ( + 0 if "adaptive_policy" not in report else report["adaptive_policy"]["episodes"] + ), + } + print(json.dumps(summary, sort_keys=True)) + return 0 + except (ImportError, RuntimeError, ValueError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/train_speed_policy.py b/scripts/train_speed_policy.py new file mode 100644 index 0000000..1365fb2 --- /dev/null +++ b/scripts/train_speed_policy.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +"""Train a Rainbow speed policy around a scripted or external chunked policy.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from experiment_config import defaults_from_argv # noqa: E402 +from policy_speed_env import make_speed_reward # noqa: E402 +from scripts.policy_cli import ( # noqa: E402 + add_base_policy_arguments, + add_observation_arguments, + build_speed_env, + comma_floats, +) +from speed_training import ( # noqa: E402 + RainbowTrainingConfig, + evaluate_rainbow_speed_policy, + train_rainbow_speed_policy, +) + + +def parse_args(): + config_defaults, config_metadata = defaults_from_argv() + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--config", + help="JSON experiment manifest or named preset such as paper-sim.", + ) + parser.add_argument("--task", choices=("pick_and_place", "insertion", "tea_bag"), default="tea_bag") + parser.add_argument("--output", type=Path, default=Path("outputs/speed_policy.pt")) + parser.add_argument( + "--report", + type=Path, + help="Optional JSON sidecar containing the summary and episode history.", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--device", default="cpu") + parser.add_argument("--speed-values", type=comma_floats, default=(1.0, 1.5, 2.0, 2.5, 3.0)) + parser.add_argument("--decisions", type=int, default=5_000) + parser.add_argument("--memory-size", type=int, default=100_000) + parser.add_argument("--batch-size", type=int, default=128) + parser.add_argument("--learning-starts", type=int, default=512) + parser.add_argument("--frame-skip", type=int, default=10) + parser.add_argument("--gradient-steps", type=int, default=4) + parser.add_argument("--hidden-dim", type=int, default=256) + parser.add_argument("--train-interval", type=int, default=1) + parser.add_argument( + "--update-schedule", + choices=("decision", "episode"), + default="decision", + help="Optimize after each decision or batch equivalent updates at episode end.", + ) + parser.add_argument( + "--checkpoint-interval", + type=int, + default=0, + help="Also save a numbered checkpoint every N decisions (0 disables it).", + ) + parser.add_argument("--target-update", type=int, default=50) + parser.add_argument("--norm-update-interval", type=int, default=100) + parser.add_argument("--learning-rate", type=float, default=1e-4) + parser.add_argument("--gamma", type=float, default=0.97) + parser.add_argument("--tau", type=float, default=0.5) + parser.add_argument("--epsilon", type=float, default=1.0) + parser.add_argument("--epsilon-decay", type=float, default=0.999) + parser.add_argument("--min-epsilon", type=float, default=0.1) + parser.add_argument("--exploration-steps", type=int, default=2000) + parser.add_argument("--per-alpha", type=float, default=0.2) + parser.add_argument("--per-beta", type=float, default=0.6) + parser.add_argument( + "--beta-schedule", + choices=("linear", "legacy"), + default="linear", + help="Importance-sampling schedule; legacy matches the retained trainer.", + ) + parser.add_argument("--atom-size", type=int, default=121) + parser.add_argument("--v-min", type=float, default=0.0) + parser.add_argument("--v-max", type=float, default=120.0) + parser.add_argument("--n-step", type=int, default=3) + parser.add_argument("--success-bonus", type=float, default=100.0) + parser.add_argument("--speed-weight", type=float, default=0.01) + parser.add_argument("--speed-power", type=float, default=2.0) + parser.add_argument("--eval-episodes", type=int, default=0) + parser.add_argument( + "--quiet", + action="store_true", + help="Suppress per-episode progress while retaining the final JSON summary.", + ) + add_base_policy_arguments(parser) + add_observation_arguments(parser) + parser.set_defaults(**config_defaults) + return parser.parse_args(), config_metadata + + +def main(): + args, config_metadata = parse_args() + reward_fn = make_speed_reward( + success_bonus=args.success_bonus, + speed_weight=args.speed_weight, + speed_power=args.speed_power, + ) + try: + env = build_speed_env(args, reward_fn=reward_fn) + config = RainbowTrainingConfig( + decisions=args.decisions, + memory_size=args.memory_size, + batch_size=args.batch_size, + learning_starts=args.learning_starts, + frame_skip=args.frame_skip, + gradient_steps=args.gradient_steps, + hidden_dim=args.hidden_dim, + train_interval=args.train_interval, + update_schedule=args.update_schedule, + checkpoint_interval=args.checkpoint_interval, + target_update=args.target_update, + norm_update_interval=args.norm_update_interval, + learning_rate=args.learning_rate, + gamma=args.gamma, + tau=args.tau, + epsilon=args.epsilon, + epsilon_decay=args.epsilon_decay, + min_epsilon=args.min_epsilon, + exploration_steps=args.exploration_steps, + alpha=args.per_alpha, + beta=args.per_beta, + beta_schedule=args.beta_schedule, + atom_size=args.atom_size, + v_min=args.v_min, + v_max=args.v_max, + n_step=args.n_step, + ) + result = train_rainbow_speed_policy( + env, + args.output, + config=config, + seed=args.seed, + device=args.device, + progress=not args.quiet, + metadata={ + "task": args.task, + "base_policy": args.base_policy, + "randomize_object_pose": args.randomize_object_pose, + "experiment_config": config_metadata, + }, + ) + summary = {key: value for key, value in result.items() if key != "episode_history"} + if args.eval_episodes: + summary["evaluation"] = evaluate_rainbow_speed_policy( + env, + args.output, + episodes=args.eval_episodes, + device=args.device, + ) + if args.report is not None: + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text( + json.dumps( + { + "summary": summary, + "episode_history": result["episode_history"], + "experiment_config": config_metadata, + }, + indent=2, + sort_keys=True, + ) + + "\n" + ) + summary["report"] = str(args.report) + print(json.dumps(summary, sort_keys=True)) + return 0 if result["losses_finite"] else 1 + except (ImportError, RuntimeError, ValueError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + finally: + if "env" in locals(): + env.close() + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/sim_env.py b/sim_env.py new file mode 100644 index 0000000..a0f6d5e --- /dev/null +++ b/sim_env.py @@ -0,0 +1,221 @@ +"""MuJoCo environments controlled by bimanual robot joint positions.""" + +from __future__ import annotations + +import collections +import os + +import numpy as np +from dm_control import mujoco +from dm_control.rl import control +from dm_control.suite import base + +from constants import ( + DT, + PUPPET_GRIPPER_POSITION_NORMALIZE_FN, + PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN, + PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN, + START_ARM_POSE, + XML_DIR, +) +from sim_tasks import ( + get_task_spec, + insertion_reward, + normalize_task_name, + sample_box_pose, + sample_insertion_pose, + tea_bag_reward, + transfer_cube_reward, +) + + +# Historical replay scripts set this after constructing the environment and before +# reset. New code may pass object_pose directly to make_sim_env. +BOX_POSE = [None] + + +def make_sim_env( + task_name: str, + render_images: bool = True, + seed: int | None = None, + object_pose=None, + randomize_object_pose: bool = False, +): + """Create a joint-control environment for any reconstructed sim task.""" + + task_name = normalize_task_name(task_name) + spec = get_task_spec(task_name) + physics = mujoco.Physics.from_xml_path(os.path.join(XML_DIR, spec.joint_xml)) + random_state = np.random.RandomState(seed) + task_classes = { + "pick_and_place": TransferCubeTask, + "insertion": InsertionTask, + "tea_bag": TransferTeaBagTask, + } + task = task_classes[task_name]( + random=random_state, + render_images=render_images, + object_pose=object_pose, + randomize_object_pose=randomize_object_pose, + ) + return control.Environment( + physics, + task, + time_limit=20, + control_timestep=DT, + n_sub_steps=None, + flat_observation=False, + ) + + +class BimanualViperXTask(base.Task): + def __init__( + self, + random=None, + render_images: bool = True, + object_pose=None, + randomize_object_pose: bool = False, + ): + super().__init__(random=random) + self.render_images = render_images + self.object_pose = None if object_pose is None else np.asarray(object_pose).copy() + self.randomize_object_pose = bool(randomize_object_pose) + + def before_step(self, action, physics): + action = np.asarray(action, dtype=np.float64) + if action.shape != (14,) or not np.all(np.isfinite(action)): + raise ValueError("Joint-control actions must be a finite array with shape (14,)") + left_gripper = PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN(action[6]) + right_gripper = PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN(action[13]) + env_action = np.concatenate( + [ + action[:6], + [left_gripper, -left_gripper], + action[7:13], + [right_gripper, -right_gripper], + ] + ) + super().before_step(env_action, physics) + + @staticmethod + def get_qpos(physics): + qpos = physics.data.qpos.copy() + left, right = qpos[:8], qpos[8:16] + return np.concatenate( + [ + left[:6], + [PUPPET_GRIPPER_POSITION_NORMALIZE_FN(left[6])], + right[:6], + [PUPPET_GRIPPER_POSITION_NORMALIZE_FN(right[6])], + ] + ) + + @staticmethod + def get_qvel(physics): + qvel = physics.data.qvel.copy() + left, right = qvel[:8], qvel[8:16] + return np.concatenate( + [ + left[:6], + [PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN(left[6])], + right[:6], + [PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN(right[6])], + ] + ) + + @staticmethod + def get_env_state(physics): + return physics.data.qpos.copy()[16:] + + def get_observation(self, physics): + obs = collections.OrderedDict( + qpos=self.get_qpos(physics), + qvel=self.get_qvel(physics), + env_state=self.get_env_state(physics), + ) + if self.render_images: + obs["images"] = { + "top": physics.render(height=480, width=640, camera_id="top"), + "angle": physics.render(height=480, width=640, camera_id="angle"), + "vis": physics.render(height=480, width=640, camera_id="front_close"), + } + return obs + + def _requested_pose(self): + if self.object_pose is not None: + return self.object_pose + if BOX_POSE[0] is not None: + return np.asarray(BOX_POSE[0]) + return None + + def _initialize_robot(self, physics): + physics.named.data.qpos[:16] = START_ARM_POSE + np.copyto(physics.data.ctrl, START_ARM_POSE) + + +class TransferCubeTask(BimanualViperXTask): + max_reward = 4 + + def initialize_episode(self, physics): + with physics.reset_context(): + self._initialize_robot(physics) + pose = self._requested_pose() + physics.named.data.qpos["red_box_joint"] = ( + sample_box_pose(self.random) if pose is None else pose + ) + super().initialize_episode(physics) + + def get_reward(self, physics): + return transfer_cube_reward(physics) + + +class InsertionTask(BimanualViperXTask): + max_reward = 4 + + def initialize_episode(self, physics): + with physics.reset_context(): + self._initialize_robot(physics) + pose = self._requested_pose() + if pose is None: + peg_pose, socket_pose = sample_insertion_pose(self.random) + pose = np.concatenate([peg_pose, socket_pose]) + if np.asarray(pose).shape != (14,): + raise ValueError("Insertion object_pose must have shape (14,)") + physics.named.data.qpos["red_peg_joint"] = pose[:7] + physics.named.data.qpos["blue_socket_joint"] = pose[7:] + super().initialize_episode(physics) + + def get_reward(self, physics): + return insertion_reward(physics) + + +class TransferTeaBagTask(BimanualViperXTask): + max_reward = 3 + + def initialize_episode(self, physics): + with physics.reset_context(): + self._initialize_robot(physics) + pose = self._requested_pose() + if pose is None and self.randomize_object_pose: + physics.named.data.qpos["red_box_joint"] = sample_box_pose(self.random) + elif pose is None: + physics.named.data.qpos["red_box_joint"] = [ + 0.15, + 0.5, + 0.05, + 1, + 0, + 0, + 0, + ] + else: + pose = np.asarray(pose) + if pose.shape != physics.data.qpos[16:].shape: + raise ValueError( + f"Tea-bag object_pose must have shape {physics.data.qpos[16:].shape}" + ) + physics.data.qpos[16:] = pose + super().initialize_episode(physics) + + def get_reward(self, physics): + return tea_bag_reward(physics) diff --git a/sim_tasks.py b/sim_tasks.py new file mode 100644 index 0000000..4d3b2ff --- /dev/null +++ b/sim_tasks.py @@ -0,0 +1,181 @@ +"""Shared definitions for the reconstructed SpeedTuning simulation tasks.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Iterable + +import numpy as np + + +@dataclass(frozen=True) +class TaskSpec: + """Metadata shared by the end-effector and joint-control environments.""" + + name: str + legacy_name: str + episode_len: int + ee_xml: str + joint_xml: str + + +TASK_SPECS = { + "pick_and_place": TaskSpec( + name="pick_and_place", + legacy_name="sim_transfer_cube_scripted", + episode_len=400, + ee_xml="bimanual_viperx_ee_transfer_cube.xml", + joint_xml="bimanual_viperx_transfer_cube.xml", + ), + "insertion": TaskSpec( + name="insertion", + legacy_name="sim_insertion_scripted", + episode_len=400, + ee_xml="bimanual_viperx_ee_insertion.xml", + joint_xml="bimanual_viperx_insertion.xml", + ), + "tea_bag": TaskSpec( + name="tea_bag", + legacy_name="sim_transfer_tea_bag_scripted", + episode_len=500, + ee_xml="bimanual_viperx_ee_transfer_tea_bag.xml", + joint_xml="bimanual_viperx_transfer_tea_bag.xml", + ), +} + +_ALIASES = { + "pick_and_place": "pick_and_place", + "pick_place": "pick_and_place", + "transfer_cube": "pick_and_place", + "sim_transfer_cube": "pick_and_place", + "sim_transfer_cube_scripted": "pick_and_place", + "sim_transfer_cube_human": "pick_and_place", + "insertion": "insertion", + "sim_insertion": "insertion", + "sim_insertion_scripted": "insertion", + "sim_insertion_human": "insertion", + "tea_bag": "tea_bag", + "teabag": "tea_bag", + "transfer_tea_bag": "tea_bag", + "sim_transfer_tea_bag": "tea_bag", + "sim_transfer_tea_bag_scripted": "tea_bag", +} + + +def normalize_task_name(task_name: str) -> str: + """Return the public task name while accepting names used by the old code.""" + + normalized = task_name.strip().lower().replace("-", "_").replace(" ", "_") + try: + return _ALIASES[normalized] + except KeyError as exc: + supported = ", ".join(TASK_SPECS) + raise ValueError(f"Unknown task {task_name!r}. Supported tasks: {supported}") from exc + + +def get_task_spec(task_name: str) -> TaskSpec: + return TASK_SPECS[normalize_task_name(task_name)] + + +def sample_box_pose(random_state=None) -> np.ndarray: + """Sample the cube pose from the range used by the original ACT simulator.""" + + rng = np.random if random_state is None else random_state + position = rng.uniform([0.0, 0.4, 0.05], [0.2, 0.6, 0.05]) + return np.concatenate([position, [1.0, 0.0, 0.0, 0.0]]) + + +def sample_insertion_pose(random_state=None) -> tuple[np.ndarray, np.ndarray]: + """Sample peg and socket poses from the original simulator ranges.""" + + rng = np.random if random_state is None else random_state + peg_position = rng.uniform([0.1, 0.4, 0.05], [0.2, 0.6, 0.05]) + socket_position = rng.uniform([-0.2, 0.4, 0.05], [-0.1, 0.6, 0.05]) + identity_quaternion = [1.0, 0.0, 0.0, 0.0] + return ( + np.concatenate([peg_position, identity_quaternion]), + np.concatenate([socket_position, identity_quaternion]), + ) + + +def contact_pairs(physics) -> set[frozenset[str]]: + """Collect active contacts without depending on MuJoCo's geom ordering.""" + + pairs = set() + for contact in physics.data.contact[: physics.data.ncon]: + geom_1 = physics.model.id2name(contact.geom1, "geom") + geom_2 = physics.model.id2name(contact.geom2, "geom") + if geom_1 is not None and geom_2 is not None: + pairs.add(frozenset((geom_1, geom_2))) + return pairs + + +def _touching(pairs: set[frozenset[str]], first: str, second: str) -> bool: + return frozenset((first, second)) in pairs + + +def _gripper_touch( + pairs: set[frozenset[str]], object_geom: str, side: str +) -> bool: + return any( + _touching(pairs, object_geom, f"vx300s_{side}/10_{finger}_gripper_finger") + for finger in ("left", "right") + ) + + +def transfer_cube_reward(physics) -> int: + pairs = contact_pairs(physics) + left = _gripper_touch(pairs, "red_box", "left") + right = _gripper_touch(pairs, "red_box", "right") + table = _touching(pairs, "red_box", "table") + + if left and not table: + return 4 + if left: + return 3 + if right and not table: + return 2 + if right: + return 1 + return 0 + + +def insertion_reward(physics) -> int: + pairs = contact_pairs(physics) + socket_geoms: Iterable[str] = ("socket-1", "socket-2", "socket-3", "socket-4") + right = _gripper_touch(pairs, "red_peg", "right") + left = any( + _gripper_touch(pairs, socket, "left") + for socket in socket_geoms + ) + peg_on_table = _touching(pairs, "red_peg", "table") + socket_on_table = any(_touching(pairs, socket, "table") for socket in socket_geoms) + peg_in_socket = any(_touching(pairs, "red_peg", socket) for socket in socket_geoms) + pin_touched = _touching(pairs, "red_peg", "pin") + + if pin_touched: + return 4 + if peg_in_socket and not peg_on_table and not socket_on_table: + return 3 + if left and right and not peg_on_table and not socket_on_table: + return 2 + if left and right: + return 1 + return 0 + + +def tea_bag_reward(physics) -> int: + pairs = contact_pairs(physics) + right = _gripper_touch(pairs, "red_box", "right") + on_table = _touching(pairs, "tea_bag", "table") or _touching( + pairs, "red_box", "table" + ) + in_cup = _touching(pairs, "cup_base", "tea_bag") + + if in_cup: + return 3 + if right and not on_table: + return 2 + if right: + return 1 + return 0 diff --git a/speed_evaluation.py b/speed_evaluation.py new file mode 100644 index 0000000..4b4d84e --- /dev/null +++ b/speed_evaluation.py @@ -0,0 +1,93 @@ +"""Deterministic evaluation and plotting helpers for SpeedTuning experiments.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np + +from speed_policy import FixedSpeedPolicy, rollout_speed_policy, summarize_rollouts + + +def speed_grid(start, stop, step): + start, stop, step = float(start), float(stop), float(step) + if step <= 0 or stop < start: + raise ValueError("Speed grid requires step > 0 and stop >= start") + count = int(np.floor((stop - start) / step + 1e-9)) + 1 + values = start + np.arange(count, dtype=np.float64) * step + if values[-1] < stop - 1e-9: + values = np.append(values, stop) + return tuple(float(np.round(value, 10)) for value in values) + + +def evaluate_seeded_policy(env_factory, policy, seeds, frame_skip=10): + rollouts = [] + for seed in seeds: + env = env_factory(int(seed)) + try: + result = rollout_speed_policy( + env, + policy, + frame_skip=frame_skip, + ) + result["seed"] = int(seed) + rollouts.append(result) + finally: + env.close() + return {**summarize_rollouts(rollouts), "seeds": list(seeds), "rollouts": rollouts} + + +def evaluate_fixed_speed_sweep(env_factory, speeds, seeds, frame_skip=10): + points = [] + for speed in speeds: + result = evaluate_seeded_policy( + env_factory, + FixedSpeedPolicy(speed), + seeds, + frame_skip=frame_skip, + ) + points.append({"fixed_speed": float(speed), **result}) + return points + + +def plot_speed_success_tradeoff(report, output_path): + """Write a compact success-versus-physical-acceleration plot.""" + + try: + import matplotlib.pyplot as plt + except ImportError as exc: + raise RuntimeError("Plotting requires: uv sync --extra evaluation") from exc + + baseline = report.get("fixed_speed_sweep", []) + if not baseline: + raise ValueError("Report has no fixed-speed sweep to plot") + figure, axis = plt.subplots(figsize=(4.4, 3.2)) + axis.plot( + [point["mean_acceleration"] for point in baseline], + [point["success_rate"] for point in baseline], + marker="o", + markersize=3, + linewidth=1.2, + label="Fixed speed", + ) + adaptive = report.get("adaptive_policy") + if adaptive is not None: + axis.scatter( + [adaptive["mean_acceleration"]], + [adaptive["success_rate"]], + marker="*", + s=90, + label="Speed policy", + zorder=3, + ) + axis.set_xlabel("Physical acceleration") + axis.set_ylabel("Success rate") + axis.set_ylim(-0.02, 1.02) + axis.grid(alpha=0.25) + axis.legend(frameon=False) + figure.tight_layout() + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + figure.savefig(output_path, dpi=200) + plt.close(figure) + return output_path diff --git a/speed_observation.py b/speed_observation.py new file mode 100644 index 0000000..3e3b72f --- /dev/null +++ b/speed_observation.py @@ -0,0 +1,284 @@ +"""Observation encoders for SpeedTuning speed policies. + +The speed learner consumes finite one-dimensional feature vectors. Encoders in +this module turn simulator observation dictionaries into those vectors while +keeping image preprocessing independent from the task policy. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Mapping, Sequence + +import numpy as np + + +def _finite_vector(value: Any, name: str) -> np.ndarray: + vector = np.asarray(value, dtype=np.float32).reshape(-1) + if vector.size == 0 or not np.all(np.isfinite(vector)): + raise ValueError(f"{name} must produce a non-empty finite feature vector") + return vector + + +@dataclass +class StateObservationEncoder: + """Encode proprioception and, optionally, privileged simulator state.""" + + include_qpos: bool = True + include_qvel: bool = True + include_env_state: bool = True + requires_images: bool = False + + def __post_init__(self): + if not (self.include_qpos or self.include_qvel or self.include_env_state): + raise ValueError("At least one state observation field must be enabled") + + def reset(self): + return None + + def __call__(self, observation: Mapping[str, Any]) -> np.ndarray: + fields = [] + for enabled, name in ( + (self.include_env_state, "env_state"), + (self.include_qpos, "qpos"), + (self.include_qvel, "qvel"), + ): + if enabled: + if name not in observation: + raise ValueError(f"Observation is missing required field {name!r}") + fields.append(_finite_vector(observation[name], name)) + return np.concatenate(fields).astype(np.float32, copy=False) + + def output_dim(self, env_state_dim: int) -> int: + return ( + (14 if self.include_qpos else 0) + + (14 if self.include_qvel else 0) + + (int(env_state_dim) if self.include_env_state else 0) + ) + + def spec(self) -> dict[str, Any]: + return { + "type": "state", + "include_qpos": self.include_qpos, + "include_qvel": self.include_qvel, + "include_env_state": self.include_env_state, + } + + +class ResNet18ImageEncoder: + """Frozen torchvision ResNet-18 image features for one or more cameras.""" + + feature_dim = 512 + + def __init__( + self, + pretrained=True, + image_size=224, + device=None, + initialize_pretrained=True, + ): + try: + import torch + from torchvision.models import ResNet18_Weights, resnet18 + except ImportError as exc: + raise RuntimeError( + "Visual speed observations require: uv sync --extra learned" + ) from exc + + self.torch = torch + self.pretrained = bool(pretrained) + self.image_size = int(image_size) + if self.image_size <= 0: + raise ValueError("image_size must be positive") + self.device = torch.device( + device or ("cuda" if torch.cuda.is_available() else "cpu") + ) + weights = ( + ResNet18_Weights.DEFAULT + if self.pretrained and initialize_pretrained + else None + ) + self.model = resnet18(weights=weights) + self.model.fc = torch.nn.Identity() + self.model.to(self.device).eval() + for parameter in self.model.parameters(): + parameter.requires_grad_(False) + self.mean = torch.tensor( + [0.485, 0.456, 0.406], dtype=torch.float32, device=self.device + ).view(1, 3, 1, 1) + self.std = torch.tensor( + [0.229, 0.224, 0.225], dtype=torch.float32, device=self.device + ).view(1, 3, 1, 1) + + def __call__(self, images: np.ndarray) -> np.ndarray: + torch = self.torch + values = np.asarray(images) + if values.ndim != 4 or values.shape[-1] != 3: + raise ValueError("Camera images must have shape [camera, height, width, 3]") + tensor = torch.as_tensor(values, dtype=torch.float32, device=self.device) + tensor = tensor.permute(0, 3, 1, 2) / 255.0 + tensor = torch.nn.functional.interpolate( + tensor, + size=(self.image_size, self.image_size), + mode="bilinear", + align_corners=False, + ) + tensor = (tensor - self.mean) / self.std + with torch.inference_mode(): + features = self.model(tensor) + return features.detach().cpu().numpy().astype(np.float32, copy=False) + + def spec(self) -> dict[str, Any]: + return { + "type": "resnet18", + "pretrained": self.pretrained, + "image_size": self.image_size, + "feature_dim": self.feature_dim, + } + + def state_dict(self): + return { + key: value.detach().cpu() + for key, value in self.model.state_dict().items() + } + + def load_state_dict(self, state_dict): + self.model.load_state_dict(state_dict) + + +class VisualObservationEncoder: + """Fuse independent image embeddings with proprioceptive features.""" + + requires_images = True + + def __init__( + self, + camera_names: Sequence[str] = ("top", "angle", "vis"), + image_encoder=None, + pretrained=True, + image_size=224, + device=None, + include_qpos=True, + include_qvel=True, + include_env_state=False, + initialize_pretrained=True, + ): + self.camera_names = tuple(camera_names) + if not self.camera_names: + raise ValueError("At least one camera is required for visual observations") + self.image_encoder = image_encoder or ResNet18ImageEncoder( + pretrained=pretrained, + image_size=image_size, + device=device, + initialize_pretrained=initialize_pretrained, + ) + self.state_encoder = StateObservationEncoder( + include_qpos=include_qpos, + include_qvel=include_qvel, + include_env_state=include_env_state, + ) + + def reset(self): + reset = getattr(self.image_encoder, "reset", None) + if reset is not None: + reset() + + def __call__(self, observation: Mapping[str, Any]) -> np.ndarray: + if "images" not in observation: + raise ValueError("Visual observations require the simulator images field") + missing = [name for name in self.camera_names if name not in observation["images"]] + if missing: + raise ValueError(f"Observation is missing cameras: {', '.join(missing)}") + images = np.stack( + [np.asarray(observation["images"][name]) for name in self.camera_names] + ) + encode = getattr(self.image_encoder, "encode", None) or self.image_encoder + image_features = _finite_vector(encode(images), "image encoder") + state_features = self.state_encoder(observation) + return np.concatenate([state_features, image_features]).astype( + np.float32, copy=False + ) + + def output_dim(self, env_state_dim: int) -> int: + feature_dim = getattr(self.image_encoder, "feature_dim", None) + if feature_dim is None: + raise AttributeError("External image encoder does not declare feature_dim") + return self.state_encoder.output_dim(env_state_dim) + len(self.camera_names) * int( + feature_dim + ) + + def spec(self) -> dict[str, Any]: + image_spec = getattr(self.image_encoder, "spec", None) + if image_spec is None: + image_spec = {"type": type(self.image_encoder).__name__} + else: + image_spec = image_spec() + return { + "type": "visual", + "camera_names": list(self.camera_names), + "state": self.state_encoder.spec(), + "image_encoder": image_spec, + } + + def state_dict(self): + state_dict = getattr(self.image_encoder, "state_dict", None) + return None if state_dict is None else state_dict() + + def load_state_dict(self, state_dict): + load = getattr(self.image_encoder, "load_state_dict", None) + if load is None: + if state_dict: + raise ValueError("The configured image encoder cannot load checkpoint state") + return + load(state_dict) + + +class ObservationEncoderAdapter: + """Validate a user-supplied observation encoder without constraining its model.""" + + def __init__(self, encoder): + if not callable(encoder): + raise TypeError("An observation encoder must be callable") + self.encoder = encoder + self.requires_images = bool(getattr(encoder, "requires_images", False)) + + def reset(self): + reset = getattr(self.encoder, "reset", None) + if reset is not None: + reset() + + def __call__(self, observation): + return _finite_vector(self.encoder(observation), "observation encoder") + + def output_dim(self, env_state_dim): + output_dim = getattr(self.encoder, "output_dim", None) + if output_dim is None: + raise AttributeError("External observation encoder does not declare output_dim") + return int(output_dim(env_state_dim) if callable(output_dim) else output_dim) + + def spec(self): + spec = getattr(self.encoder, "spec", None) + return ( + {"type": type(self.encoder).__name__} + if spec is None + else dict(spec() if callable(spec) else spec) + ) + + def state_dict(self): + state_dict = getattr(self.encoder, "state_dict", None) + return None if state_dict is None else state_dict() + + def load_state_dict(self, state_dict): + load = getattr(self.encoder, "load_state_dict", None) + if load is None: + if state_dict: + raise ValueError("External observation encoder cannot load checkpoint state") + return + load(state_dict) + + +def encoder_spec(encoder) -> dict[str, Any]: + spec = getattr(encoder, "spec", None) + if spec is None: + return {"type": type(encoder).__name__} + return dict(spec() if callable(spec) else spec) diff --git a/speed_policy.py b/speed_policy.py new file mode 100644 index 0000000..05caf9e --- /dev/null +++ b/speed_policy.py @@ -0,0 +1,322 @@ +"""Model-agnostic speed-policy adapters and rollout utilities.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Sequence + +import numpy as np + + +@dataclass(frozen=True) +class SpeedContext: + """Episode metadata passed to a speed policy at each decision.""" + + policy_time: float + physics_steps: int + episode_len: int + speed_values: tuple[float, ...] + + +class SpeedPolicyAdapter: + """Adapt a callable or ``select_speed`` object to the speed contract. + + A public speed policy receives ``(observation, context)`` and returns a + finite positive multiplier such as ``1.0`` or ``1.5``. + """ + + def __init__(self, policy: Any): + select = getattr(policy, "select_speed", None) + if select is None and callable(policy): + select = policy + if select is None: + raise TypeError("speed policy must be callable or define select_speed()") + self.policy = policy + self._select = select + + def reset(self): + reset = getattr(self.policy, "reset", None) + if reset is not None: + reset() + + def __call__(self, observation, context): + speed = float(self._select(observation, context)) + if not np.isfinite(speed) or speed <= 0: + raise ValueError("A speed policy must return a finite positive multiplier") + return speed + + +class FixedSpeedPolicy: + """Always choose one physical speed multiplier.""" + + def __init__(self, speed=1.0): + self.speed = float(speed) + if not np.isfinite(self.speed) or self.speed <= 0: + raise ValueError("speed must be finite and positive") + + def select_speed(self, observation, context): + del observation, context + return self.speed + + +class SpeedProfilePolicy: + """Select piecewise-constant speeds over normalized nominal policy time.""" + + def __init__(self, speeds: Sequence[float]): + values = np.asarray(speeds, dtype=np.float64) + if values.ndim != 1 or len(values) == 0: + raise ValueError("speeds must be a non-empty one-dimensional sequence") + if not np.all(np.isfinite(values)) or np.any(values <= 0): + raise ValueError("Every profile speed must be finite and positive") + self.speeds = tuple(float(value) for value in values) + + def select_speed(self, observation, context): + del observation + fraction = min(max(context.policy_time / context.episode_len, 0.0), 1.0) + index = min(int(fraction * len(self.speeds)), len(self.speeds) - 1) + return self.speeds[index] + + +class CallableSpeedPolicy: + """Give a descriptive wrapper to a user-provided speed function.""" + + def __init__(self, function: Callable[[np.ndarray, SpeedContext], float]): + self.function = function + + def select_speed(self, observation, context): + return self.function(observation, context) + + +class RainbowSpeedPolicy: + """Inference-only speed policy loaded from a public training checkpoint.""" + + def __init__( + self, + network, + speed_values, + device="cpu", + frame_skip=10, + observation_spec=None, + environment_spec=None, + observation_encoder_state_dict=None, + checkpoint_metadata=None, + ): + import torch + + self.torch = torch + self.device = torch.device(device) + self.network = network.to(self.device).eval() + self.speed_values = tuple(float(value) for value in speed_values) + self.observation_dim = int(network.in_dim) + self.frame_skip = int(frame_skip) + self.observation_spec = observation_spec + self.environment_spec = environment_spec + self.observation_encoder_state_dict = observation_encoder_state_dict + self.checkpoint_metadata = dict(checkpoint_metadata or {}) + + @classmethod + def load(cls, checkpoint_path, device="cpu"): + try: + import torch + except ImportError as exc: + raise RuntimeError("Rainbow evaluation requires: uv sync --extra rl") from exc + from rl.rainbowDQN.network import Network + + checkpoint_path = Path(checkpoint_path) + # Public speed checkpoints contain tensors and primitive metadata only, + # so use PyTorch's restricted unpickler for downloaded artifacts. + payload = torch.load( + checkpoint_path, map_location=device, weights_only=True + ) + required = { + "model_state_dict", + "observation_dim", + "speed_values", + "atom_size", + "v_min", + "v_max", + "hidden_dim", + } + missing = sorted(required.difference(payload)) + if missing: + raise ValueError(f"Speed checkpoint is missing keys: {', '.join(missing)}") + support = torch.linspace( + float(payload["v_min"]), + float(payload["v_max"]), + int(payload["atom_size"]), + ).to(device) + network = Network( + int(payload["observation_dim"]), + len(payload["speed_values"]), + int(payload["atom_size"]), + support, + hidden_dim=int(payload["hidden_dim"]), + ) + network.load_state_dict(payload["model_state_dict"]) + training_config = payload.get("training_config", {}) + return cls( + network, + payload["speed_values"], + device=device, + frame_skip=payload.get( + "decision_frame_skip", training_config.get("frame_skip", 10) + ), + observation_spec=payload.get("observation_spec"), + environment_spec=payload.get("environment_spec"), + observation_encoder_state_dict=payload.get( + "observation_encoder_state_dict" + ), + checkpoint_metadata=payload.get("metadata"), + ) + + def select_action(self, observation): + tensor = self.torch.as_tensor( + observation, dtype=self.torch.float32, device=self.device + ).unsqueeze(0) + with self.torch.inference_mode(): + action = int(self.network(tensor).argmax(dim=1).item()) + return action + + def select_speed(self, observation, context): + del context + return self.speed_values[self.select_action(observation)] + + def configure_environment(self, env): + """Restore and validate observation preprocessing for evaluation.""" + + env.load_observation_encoder_state_dict( + self.observation_encoder_state_dict + ) + observation = env.reset() + if self.observation_dim != observation.size: + raise ValueError( + "Checkpoint observation size does not match the environment: " + f"{self.observation_dim} != {observation.size}" + ) + if tuple(self.speed_values) != tuple(env.speed_values): + raise ValueError( + "Checkpoint speed_values do not match the environment: " + f"{self.speed_values} != {env.speed_values}" + ) + if self.observation_spec is not None: + actual_spec = env.observation_spec() + if self.observation_spec != actual_spec: + raise ValueError( + "Checkpoint observation preprocessing does not match the environment: " + f"{self.observation_spec!r} != {actual_spec!r}" + ) + if self.environment_spec is not None: + actual_environment = env.environment_spec() + if self.environment_spec != actual_environment: + raise ValueError( + "Checkpoint environment does not match evaluation: " + f"{self.environment_spec!r} != {actual_environment!r}" + ) + return observation + + +def rollout_speed_policy(env, speed_policy, capture_speeds=False, frame_skip=None): + """Pair any conforming speed policy with a configured speed environment.""" + + raw_policy = speed_policy + prepared_observation = None + if isinstance(raw_policy, RainbowSpeedPolicy): + prepared_observation = raw_policy.configure_environment(env) + policy = ( + speed_policy + if isinstance(speed_policy, SpeedPolicyAdapter) + else SpeedPolicyAdapter(speed_policy) + ) + policy.reset() + observation = ( + env.reset() if prepared_observation is None else prepared_observation + ) + decision_frame_skip = int( + frame_skip + if frame_skip is not None + else getattr(raw_policy, "frame_skip", env.decision_frame_skip) + ) + if decision_frame_skip <= 0: + raise ValueError("frame_skip must be positive") + done = False + total_reward = 0.0 + info = {"success": False} + speeds = [] + decisions = 0 + while not done: + context = SpeedContext( + policy_time=env.policy_time, + physics_steps=env.physics_steps, + episode_len=env.episode_len, + speed_values=env.speed_values, + ) + speed = policy(observation, context) + observation, reward, done, info = env.step_decision( + speed, + frame_skip=decision_frame_skip, + quantized=False, + ) + total_reward += reward + decisions += 1 + if capture_speeds: + speeds.append(speed) + + acceleration = float(env.episode_len / max(info["physics_steps"], 1)) + result = { + "success": bool(info["success"]), + "return": float(total_reward), + "physics_steps": int(info["physics_steps"]), + "policy_time": float(info["policy_time"]), + "mean_speed": float(np.mean(env.speed_list)), + "max_speed": float(np.max(env.speed_list)), + "acceleration": acceleration, + "successful_acceleration": acceleration if info["success"] else None, + "decisions": decisions, + "decision_frame_skip": decision_frame_skip, + "duration_seconds": float(info["physics_steps"] / 50.0), + "nominal_duration_seconds": float(env.episode_len / 50.0), + } + if capture_speeds: + result["speeds"] = speeds + return result + + +def summarize_rollouts(rollouts): + """Aggregate paper-facing success and physical-acceleration metrics.""" + + rollouts = list(rollouts) + if not rollouts: + raise ValueError("At least one rollout is required") + successes = np.asarray([item["success"] for item in rollouts], dtype=np.float64) + accelerations = np.asarray( + [item["acceleration"] for item in rollouts], dtype=np.float64 + ) + successful = accelerations[successes.astype(bool)] + physics_steps = np.asarray( + [item["physics_steps"] for item in rollouts], dtype=np.float64 + ) + mean_speeds = np.asarray( + [item["mean_speed"] for item in rollouts], dtype=np.float64 + ) + return { + "episodes": len(rollouts), + "successes": int(successes.sum()), + "success_rate": float(successes.mean()), + "success_standard_error": float( + np.sqrt(successes.mean() * (1.0 - successes.mean()) / len(successes)) + ), + "mean_acceleration": float(accelerations.mean()), + "median_acceleration": float(np.median(accelerations)), + "acceleration_standard_deviation": float(accelerations.std()), + "acceleration_25th_percentile": float(np.percentile(accelerations, 25)), + "acceleration_75th_percentile": float(np.percentile(accelerations, 75)), + "mean_successful_acceleration": ( + None if successful.size == 0 else float(successful.mean()) + ), + "successful_acceleration_standard_deviation": ( + None if successful.size == 0 else float(successful.std()) + ), + "mean_physics_steps": float(physics_steps.mean()), + "mean_commanded_speed": float(mean_speeds.mean()), + } diff --git a/speed_training.py b/speed_training.py new file mode 100644 index 0000000..1c02dd5 --- /dev/null +++ b/speed_training.py @@ -0,0 +1,298 @@ +"""Training and evaluation loops for the included Rainbow speed policy.""" + +from __future__ import annotations + +import random +from collections import deque +from dataclasses import asdict, dataclass +from pathlib import Path + +import numpy as np + +from speed_policy import RainbowSpeedPolicy, rollout_speed_policy, summarize_rollouts + + +@dataclass(frozen=True) +class RainbowTrainingConfig: + """Practical defaults for speed-policy experiments, not paper reproduction.""" + + decisions: int = 5_000 + memory_size: int = 100_000 + batch_size: int = 128 + learning_starts: int = 512 + frame_skip: int = 10 + gradient_steps: int = 4 + train_interval: int = 1 + target_update: int = 50 + norm_update_interval: int = 100 + learning_rate: float = 1e-4 + gamma: float = 0.97 + tau: float = 0.5 + epsilon: float = 1.0 + epsilon_decay: float = 0.999 + min_epsilon: float = 0.1 + exploration_steps: int = 2_000 + alpha: float = 0.2 + beta: float = 0.6 + beta_schedule: str = "linear" + atom_size: int = 121 + v_min: float = 0.0 + v_max: float = 120.0 + n_step: int = 3 + hidden_dim: int = 256 + update_schedule: str = "decision" + checkpoint_interval: int = 0 + + def validate(self): + positive_ints = ( + "decisions", + "memory_size", + "batch_size", + "learning_starts", + "frame_skip", + "gradient_steps", + "train_interval", + "target_update", + "norm_update_interval", + "atom_size", + "n_step", + "hidden_dim", + ) + for name in positive_ints: + if int(getattr(self, name)) <= 0: + raise ValueError(f"{name} must be positive") + if self.memory_size < self.batch_size: + raise ValueError("memory_size must be at least batch_size") + if self.learning_starts < self.batch_size: + raise ValueError("learning_starts must be at least batch_size") + if self.atom_size < 2 or self.v_max <= self.v_min: + raise ValueError("Categorical support requires atom_size >= 2 and v_max > v_min") + if self.update_schedule not in {"decision", "episode"}: + raise ValueError("update_schedule must be 'decision' or 'episode'") + if self.beta_schedule not in {"linear", "legacy"}: + raise ValueError("beta_schedule must be 'linear' or 'legacy'") + if self.checkpoint_interval < 0: + raise ValueError("checkpoint_interval cannot be negative") + + +def _normalization_stats(states): + values = np.asarray(states, dtype=np.float32) + return { + "states_mean": values.mean(axis=0), + "states_std": np.maximum(values.std(axis=0), 1e-6), + } + + +def _checkpoint_payload(agent, env, config, seed, metadata, completed_decisions): + return { + "format_version": 2, + "algorithm": "rainbow_dqn", + "model_state_dict": agent.dqn.state_dict(), + "observation_dim": int(env.obs_space), + "speed_values": list(env.speed_values), + "atom_size": config.atom_size, + "v_min": config.v_min, + "v_max": config.v_max, + "hidden_dim": config.hidden_dim, + "seed": int(seed), + "training_config": asdict(config), + "completed_decisions": int(completed_decisions), + "decision_frame_skip": config.frame_skip, + "reward_aggregation": "undiscounted_sum_per_decision", + "observation_spec": env.observation_spec(), + "environment_spec": env.environment_spec(), + "observation_encoder_state_dict": env.observation_encoder_state_dict(), + "metric_spec": { + "acceleration": "episode_len / physics_steps", + "control_frequency_hz": 50, + }, + "metadata": dict(metadata or {}), + } + + +def _numbered_checkpoint_path(checkpoint_path, decision): + checkpoint_path = Path(checkpoint_path) + return checkpoint_path.with_name( + f"{checkpoint_path.stem}.decision-{decision}{checkpoint_path.suffix}" + ) + + +def train_rainbow_speed_policy( + env, + checkpoint_path, + config=None, + seed=0, + device=None, + metadata=None, + progress=True, +): + """Train Rainbow against any ``SpeedPolicyEnv`` and save one checkpoint.""" + + try: + import torch + except ImportError as exc: + raise RuntimeError("Rainbow training requires: uv sync --extra rl") from exc + from rl.rainbowDQN.dqnAgent import DQNAgent + + config = config or RainbowTrainingConfig() + config.validate() + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + + state = env.reset() + agent = DQNAgent( + env, + memory_size=config.memory_size, + batch_size=config.batch_size, + target_update=config.target_update, + seed=seed, + lr=config.learning_rate, + gamma=config.gamma, + tau=config.tau, + frame_skip=config.frame_skip, + epsilon=config.epsilon, + epsilon_decay=config.epsilon_decay, + min_epsilon=config.min_epsilon, + exploration_steps=config.exploration_steps, + alpha=config.alpha, + beta=config.beta, + atom_size=config.atom_size, + v_min=config.v_min, + v_max=config.v_max, + n_step=config.n_step, + hidden_dim=config.hidden_dim, + device=device, + log_dir=None, + ) + + state_history = deque(maxlen=min(config.memory_size, 100_000)) + state_history.append(state.copy()) + episode_return = 0.0 + episode_decisions = 0 + episode_index = 0 + update_count = 0 + losses = [] + episodes = [] + numbered_checkpoints = [] + + def update_network(number_of_updates): + nonlocal update_count + for _ in range(number_of_updates): + losses.append(float(agent.update_model())) + update_count += 1 + if update_count % config.target_update == 0: + agent._target_soft_update() + + for decision in range(1, config.decisions + 1): + action = agent.select_action(state) + next_state, reward, done, info = agent.step(action, config.frame_skip) + state_history.append(next_state.copy()) + episode_return += float(reward) + episode_decisions += 1 + + progress_fraction = min(decision / config.decisions, 1.0) + if config.beta_schedule == "legacy": + # Retained trainer behavior: repeatedly close the remaining gap. + agent.beta += progress_fraction * (1.0 - agent.beta) + else: + agent.beta = config.beta + progress_fraction * (1.0 - config.beta) + agent.decay_epsilon(decision) + + ready = len(agent.memory) >= max(config.batch_size, config.learning_starts) + if ( + config.update_schedule == "decision" + and ready + and decision % config.train_interval == 0 + ): + if decision % config.norm_update_interval == 0 or update_count == 0: + stats = _normalization_stats(state_history) + agent.dqn.update_norm_stats(stats) + agent.dqn_target.update_norm_stats(stats) + update_network(config.gradient_steps) + + state = next_state + if done: + episode_index += 1 + record = { + "episode": episode_index, + "decision": decision, + "return": episode_return, + "decisions": episode_decisions, + "physics_steps": int(info["physics_steps"]), + "mean_speed": float(np.mean(env.speed_list)), + "acceleration": float( + env.episode_len / max(int(info["physics_steps"]), 1) + ), + "success": bool(info["success"]), + } + episodes.append(record) + if progress: + print( + "episode={episode} decision={decision} success={success} " + "return={return:.3f} mean_speed={mean_speed:.3f}".format(**record) + ) + if config.update_schedule == "episode" and ready: + stats = _normalization_stats(state_history) + agent.dqn.update_norm_stats(stats) + agent.dqn_target.update_norm_stats(stats) + # The retained SpeedTuning trainer optimized once per decision, + # batching those updates at the end of each episode. + update_network(episode_decisions * config.gradient_steps) + state = env.reset() + state_history.append(state.copy()) + episode_return = 0.0 + episode_decisions = 0 + + if ( + config.checkpoint_interval + and decision % config.checkpoint_interval == 0 + ): + numbered_path = _numbered_checkpoint_path(checkpoint_path, decision) + numbered_path.parent.mkdir(parents=True, exist_ok=True) + torch.save( + _checkpoint_payload( + agent, env, config, seed, metadata, completed_decisions=decision + ), + numbered_path, + ) + numbered_checkpoints.append(str(numbered_path)) + + if len(state_history) >= 2: + stats = _normalization_stats(state_history) + agent.dqn.update_norm_stats(stats) + + checkpoint_path = Path(checkpoint_path) + checkpoint_path.parent.mkdir(parents=True, exist_ok=True) + payload = _checkpoint_payload( + agent, + env, + config, + seed, + metadata, + completed_decisions=config.decisions, + ) + torch.save(payload, checkpoint_path) + + finite_losses = bool(np.isfinite(losses).all()) if losses else True + return { + "checkpoint": str(checkpoint_path), + "decisions": config.decisions, + "episodes": len(episodes), + "successes": sum(int(item["success"]) for item in episodes), + "updates": update_count, + "losses_finite": finite_losses, + "loss_last": losses[-1] if losses else None, + "numbered_checkpoints": numbered_checkpoints, + "episode_history": episodes, + } + + +def evaluate_rainbow_speed_policy(env, checkpoint_path, episodes=10, device="cpu"): + """Evaluate a saved speed policy against the supplied base-policy wrapper.""" + + if episodes <= 0: + raise ValueError("episodes must be positive") + policy = RainbowSpeedPolicy.load(checkpoint_path, device=device) + results = [rollout_speed_policy(env, policy) for _ in range(episodes)] + return {**summarize_rollouts(results), "rollouts": results} diff --git a/tests/test_chunked_policies.py b/tests/test_chunked_policies.py new file mode 100644 index 0000000..b9b29ef --- /dev/null +++ b/tests/test_chunked_policies.py @@ -0,0 +1,98 @@ +import numpy as np +import pytest + +from chunked_policy import ( + ChunkPredictorAdapter, + ChunkedPolicyRunner, + TorchChunkPredictor, + as_action_chunk, + interpolate_action_chunk, + replay_recorded_chunks, +) +from sim_env import make_sim_env +from sim_tasks import TASK_SPECS + + +def test_action_chunk_interpolation_changes_execution_length(): + actions = np.arange(10 * 14, dtype=float).reshape(10, 14) + accelerated = interpolate_action_chunk(actions, speed=2.0) + slowed = interpolate_action_chunk(actions, speed=0.5) + assert accelerated.shape == (5, 14) + assert slowed.shape == (20, 14) + np.testing.assert_array_equal(accelerated, actions[::2]) + + +def test_upstream_chunk_adapter_accepts_common_batch_and_dict_output(): + actions = np.arange(4 * 14, dtype=float).reshape(1, 4, 14) + adapter = ChunkPredictorAdapter(lambda observation: {"actions": actions}) + np.testing.assert_array_equal(adapter({}), actions[0]) + np.testing.assert_array_equal(as_action_chunk(actions[0, 0]), actions[0, :1]) + + +def test_chunk_runner_accepts_online_speed_changes(): + actions = np.repeat(np.arange(10, dtype=float)[:, None], 14, axis=1) + runner = ChunkedPolicyRunner(lambda observation: actions) + samples = [ + runner.action({}, speed=speed)[0] + for speed in (1.0, 2.0, 0.5, 1.5) + ] + np.testing.assert_allclose(samples, [0.0, 1.0, 3.0, 3.5]) + + +@pytest.mark.parametrize("task_name", TASK_SPECS) +def test_recorded_action_chunks_complete_joint_task(task_name): + result = replay_recorded_chunks(task_name, chunk_size=25, seed=0) + assert result["success"] + + +@pytest.mark.learned +def test_actual_act_model_accepts_every_simulator_observation(): + torch = pytest.importorskip("torch") + pytest.importorskip("torchvision") + from policy import ACTPolicy + + config = { + "lr": 1e-4, + "num_queries": 4, + "kl_weight": 1, + "hidden_dim": 32, + "dim_feedforward": 64, + "lr_backbone": 0.0, + "backbone": "resnet18", + "enc_layers": 1, + "dec_layers": 1, + "nheads": 4, + "camera_names": ["top"], + "pretrained_backbone": False, + "device": "cpu", + } + policy = ACTPolicy(config).eval() + for task_name in TASK_SPECS: + env = make_sim_env(task_name, render_images=True, seed=0) + timestep = env.reset() + initial_qpos = timestep.observation["qpos"].copy() + predictor = TorchChunkPredictor( + policy, + ["top"], + qpos_mean=np.zeros(14), + qpos_std=np.ones(14), + action_mean=initial_qpos, + action_std=np.full(14, 1e-3), + device="cpu", + ) + chunk = predictor(timestep.observation) + assert chunk.shape == (4, 14) + assert np.isfinite(chunk).all() + timestep = env.step(chunk[0]) + assert torch.isfinite(torch.as_tensor(timestep.observation["qpos"])).all() + + +@pytest.mark.rl +def test_rainbow_dqn_optimization_poc(): + pytest.importorskip("torch") + from scripts.rainbow_poc import run_poc + + result = run_poc(seed=0, min_transitions=64, updates=8) + assert result["passed"] + assert result["successful_episodes"] >= 1 + assert result["actions_seen"] == [0, 1, 2, 3, 4] diff --git a/tests/test_paper_parity.py b/tests/test_paper_parity.py new file mode 100644 index 0000000..83ebf23 --- /dev/null +++ b/tests/test_paper_parity.py @@ -0,0 +1,241 @@ +import json +from pathlib import Path + +import numpy as np +import pytest + +from chunked_policy import ChunkedPolicyRunner +from experiment_config import load_experiment_config +from policy_speed_env import create_speed_env, make_speed_reward +from speed_evaluation import speed_grid +from speed_observation import StateObservationEncoder, VisualObservationEncoder +from speed_policy import FixedSpeedPolicy, rollout_speed_policy, summarize_rollouts + + +class CountingChunkPredictor: + def __init__(self): + self.calls = 0 + + def __call__(self, observation): + self.calls += 1 + return np.repeat(np.asarray(observation["qpos"])[None], 100, axis=0) + + +def test_chunk_runner_replans_at_each_decision_boundary(): + predictor = CountingChunkPredictor() + runner = ChunkedPolicyRunner(predictor) + observation = {"qpos": np.zeros(14)} + runner.begin_decision(observation, speed=2.0) + runner.action(observation, speed=2.0) + runner.action(observation, speed=2.0) + assert predictor.calls == 1 + runner.begin_decision(observation, speed=1.5) + assert predictor.calls == 2 + + +def test_decision_step_uses_one_fresh_chunk_and_holds_speed(): + predictor = CountingChunkPredictor() + env = create_speed_env( + "tea_bag", + chunk_predictor=predictor, + render_images=False, + seed=0, + ) + env.reset() + _, _, done, info = env.step_decision(1.5, frame_skip=4, quantized=False) + assert not done + assert predictor.calls == 1 + assert info["decision_physics_steps"] == 4 + assert env.speed_list == [1.5] * 4 + env.step_decision(1.0, frame_skip=2, quantized=False) + assert predictor.calls == 2 + + +def test_stacked_unprivileged_proprioceptive_observations(): + env = create_speed_env( + "pick_and_place", + observation_encoder=StateObservationEncoder(include_env_state=False), + frame_stack=5, + seed=0, + ) + observation = env.reset() + assert observation.shape == (5 * 28,) + frames = observation.reshape(5, 28) + np.testing.assert_allclose(frames, np.repeat(frames[:1], 5, axis=0)) + next_observation, _, _, _ = env.step_decision( + 1.0, frame_skip=1, quantized=False + ) + assert next_observation.shape == observation.shape + np.testing.assert_allclose(next_observation[: 4 * 28], observation[28:]) + + +class DummyImageEncoder: + feature_dim = 2 + + def __call__(self, images): + means = images.mean(axis=(1, 2, 3)) + return np.stack([means, means + 1.0], axis=1) + + def spec(self): + return {"type": "dummy", "feature_dim": self.feature_dim} + + +def test_visual_encoder_fuses_camera_features_without_env_state(): + encoder = VisualObservationEncoder( + camera_names=("top", "angle"), + image_encoder=DummyImageEncoder(), + include_env_state=False, + ) + observation = { + "qpos": np.zeros(14), + "qvel": np.ones(14), + "env_state": np.full(7, 9.0), + "images": { + "top": np.zeros((4, 5, 3), dtype=np.uint8), + "angle": np.full((4, 5, 3), 10, dtype=np.uint8), + }, + } + features = encoder(observation) + assert features.shape == (32,) + assert not np.any(features == 9.0) + assert encoder.spec()["image_encoder"]["type"] == "dummy" + + +def test_paper_manifest_and_ablation_inheritance(): + paper, _ = load_experiment_config("paper-sim") + assert paper["speed_values"] == [1.5, 2.0, 3.0, 4.5] + assert paper["frame_stack"] == 5 + assert paper["include_env_state"] is False + assert paper["speed_weight"] == 0.01 + assert paper["update_schedule"] == "episode" + no_image_path = Path(__file__).parents[1] / "configs/ablations/no_image.json" + no_image, _ = load_experiment_config(no_image_path) + assert no_image["speed_observation"] == "state" + assert no_image["hidden_dim"] == 1024 + + scripted, _ = load_experiment_config("scripted-tea-bag") + assert scripted["base_policy"] == "scripted" + assert scripted["speed_values"] == [1.0, 1.5, 2.0, 2.5, 3.0] + assert scripted["decisions"] == 100_000 + assert scripted["learning_starts"] == 128 + assert scripted["beta_schedule"] == "legacy" + randomized, _ = load_experiment_config("scripted-tea-bag-randomized") + assert randomized["randomize_object_pose"] is True + assert randomized["update_schedule"] == "episode" + + pick, _ = load_experiment_config("scripted-pick-and-place") + insertion, _ = load_experiment_config("scripted-insertion") + assert pick["speed_values"][-1] == 4.5 + assert insertion["speed_values"] == [1.0, 1.5, 2.0, 2.5, 3.0] + assert pick["decisions"] == insertion["decisions"] == 100_000 + + +def test_paper_metrics_use_physical_length_not_commanded_speed(): + env = create_speed_env("tea_bag", seed=0, decision_frame_skip=10) + rollout = rollout_speed_policy(env, FixedSpeedPolicy(1.5)) + assert rollout["acceleration"] == pytest.approx( + env.episode_len / rollout["physics_steps"] + ) + summary = summarize_rollouts([rollout]) + assert summary["mean_acceleration"] == rollout["acceleration"] + assert "mean_commanded_speed" in summary + + +def test_speed_grid_is_stable_at_decimal_intervals(): + assert speed_grid(1.0, 1.3, 0.1) == (1.0, 1.1, 1.2, 1.3) + + +def test_zero_degree_speed_reward_ablation_is_supported(): + reward = make_speed_reward(speed_weight=0.0, speed_power=0.0) + assert reward(4.5, done=False, success=False) == 0.0 + + +def test_reference_results_cover_all_tasks_without_checkpoint_paths(): + path = Path(__file__).parents[1] / "benchmarks/scripted_results.json" + benchmark = json.loads(path.read_text()) + assert set(benchmark["results"]) == {"pick_and_place", "insertion", "tea_bag"} + serialized = json.dumps(benchmark) + assert not any(suffix in serialized for suffix in (".pt\"", ".pth\"", ".ckpt\"")) + for result in benchmark["results"].values(): + assert 0.0 <= result["learned_speed"]["success_rate"] <= 1.0 + assert result["learned_speed"]["mean_physical_acceleration"] > 1.0 + + +@pytest.mark.learned +def test_resnet_visual_speed_observation_and_state_round_trip(): + pytest.importorskip("torch") + pytest.importorskip("torchvision") + first = VisualObservationEncoder( + camera_names=("top",), + pretrained=False, + image_size=64, + device="cpu", + include_env_state=False, + ) + env = create_speed_env( + "tea_bag", + observation_encoder=first, + frame_stack=2, + render_images=True, + seed=0, + ) + observation = env.reset() + assert observation.shape == (2 * (28 + 512),) + saved_state = first.state_dict() + second = VisualObservationEncoder( + camera_names=("top",), + pretrained=False, + image_size=64, + device="cpu", + include_env_state=False, + ) + second.load_state_dict(saved_state) + assert first.spec() == second.spec() + env.close() + + +@pytest.mark.learned +def test_retained_act_checkpoint_loader_runs_on_simulator(tmp_path: Path): + torch = pytest.importorskip("torch") + pytest.importorskip("torchvision") + from act_integration import build_act_chunk_predictor + from policy import ACTPolicy + from sim_env import make_sim_env + + config = { + "lr": 1e-4, + "num_queries": 4, + "kl_weight": 1, + "hidden_dim": 32, + "dim_feedforward": 64, + "lr_backbone": 0.0, + "backbone": "resnet18", + "enc_layers": 1, + "dec_layers": 1, + "nheads": 4, + "camera_names": ["top"], + "pretrained_backbone": False, + "device": "cpu", + } + policy = ACTPolicy(config).eval() + checkpoint = tmp_path / "act.pt" + torch.save( + { + "model_state_dict": policy.state_dict(), + "policy_config": config, + "stats": { + "qpos_mean": np.zeros(14), + "qpos_std": np.ones(14), + "action_mean": np.zeros(14), + "action_std": np.ones(14), + }, + }, + checkpoint, + ) + predictor = build_act_chunk_predictor( + "tea_bag", checkpoint=checkpoint, device="cpu" + ) + env = make_sim_env("tea_bag", render_images=True, seed=0) + chunk = predictor(env.reset().observation) + assert chunk.shape == (4, 14) + assert np.isfinite(chunk).all() diff --git a/tests/test_sim_envs.py b/tests/test_sim_envs.py new file mode 100644 index 0000000..497d822 --- /dev/null +++ b/tests/test_sim_envs.py @@ -0,0 +1,107 @@ +import numpy as np +import pytest + +from ee_sim_env import make_ee_sim_env +from policy_speed_env import create_speed_env +from scripts.run_sim import run_task +from scripted_policy import make_scripted_policy +from sim_env import BOX_POSE, make_sim_env +from sim_tasks import TASK_SPECS, normalize_task_name + + +@pytest.mark.parametrize( + ("alias", "expected"), + [ + ("pick-and-place", "pick_and_place"), + ("sim_transfer_cube_scripted", "pick_and_place"), + ("sim_insertion", "insertion"), + ("teabag", "tea_bag"), + ("sim_transfer_tea_bag_scripted", "tea_bag"), + ], +) +def test_task_aliases(alias, expected): + assert normalize_task_name(alias) == expected + + +@pytest.mark.parametrize("task_name", TASK_SPECS) +def test_end_effector_env_resets_and_steps(task_name): + env = make_ee_sim_env(task_name, render_images=False, seed=7) + timestep = env.reset() + assert timestep.observation["qpos"].shape == (14,) + assert "images" not in timestep.observation + action = make_scripted_policy(task_name)(timestep) + timestep = env.step(action) + assert np.isfinite(timestep.observation["qpos"]).all() + + +@pytest.mark.parametrize("task_name", TASK_SPECS) +def test_joint_env_resets_without_external_object_pose(task_name): + BOX_POSE[0] = None + env = make_sim_env(task_name, render_images=False, seed=7) + timestep = env.reset() + timestep = env.step(timestep.observation["qpos"]) + assert timestep.observation["qpos"].shape == (14,) + assert np.isfinite(timestep.observation["env_state"]).all() + + +def test_tea_bag_pose_randomization_is_opt_in_and_seeded(): + fixed = make_ee_sim_env("tea_bag", render_images=False, seed=7) + fixed_pose = fixed.reset().observation["env_state"][:7].copy() + np.testing.assert_allclose( + fixed.reset().observation["env_state"][:7], fixed_pose + ) + + randomized = make_ee_sim_env( + "tea_bag", + render_images=False, + seed=7, + randomize_object_pose=True, + ) + first_pose = randomized.reset().observation["env_state"][:7].copy() + second_pose = randomized.reset().observation["env_state"][:7].copy() + assert not np.allclose(first_pose[:2], second_pose[:2]) + assert np.all(first_pose[:2] >= [0.0, 0.4]) + assert np.all(first_pose[:2] <= [0.2, 0.6]) + + +@pytest.mark.parametrize("task_name", TASK_SPECS) +def test_original_scripted_policy_completes_task(task_name): + spec = TASK_SPECS[task_name] + env = make_ee_sim_env(task_name, render_images=False, seed=0) + timestep = env.reset() + policy = make_scripted_policy(task_name) + rewards = [] + for _ in range(spec.episode_len): + timestep = env.step(policy(timestep)) + rewards.append(timestep.reward) + assert max(rewards) == env.task.max_reward + + +@pytest.mark.parametrize("task_name", TASK_SPECS) +def test_speed_wrapper_completes_task_at_original_speed(task_name): + env = create_speed_env(task_name=task_name, seed=0) + observation = env.reset() + assert observation.shape == (env.obs_space,) + done = False + info = {} + while not done: + observation, _, done, info = env.step(1.0, quantized=False) + assert info["success"] + + +@pytest.mark.parametrize("task_name", TASK_SPECS) +def test_scripted_policy_completes_task_at_1_5x(task_name): + result = run_task(task_name, speed=1.5, seed=0) + assert result["success"] + assert result["steps"] < TASK_SPECS[task_name].episode_len + + +@pytest.mark.parametrize("task_name", TASK_SPECS) +def test_quantized_speed_wrapper_completes_task_at_1_5x(task_name): + env = create_speed_env(task_name=task_name, seed=0) + env.reset() + done = False + info = {} + while not done: + _, _, done, info = env.step(1, quantized=True) + assert info["success"] diff --git a/tests/test_speed_infrastructure.py b/tests/test_speed_infrastructure.py new file mode 100644 index 0000000..d06298f --- /dev/null +++ b/tests/test_speed_infrastructure.py @@ -0,0 +1,158 @@ +from pathlib import Path + +import numpy as np +import pytest +from dm_control.rl import control + +from policy_speed_env import create_recorded_chunk_speed_env, create_speed_env +from speed_policy import ( + FixedSpeedPolicy, + SpeedContext, + SpeedPolicyAdapter, + SpeedProfilePolicy, + rollout_speed_policy, +) + + +def test_speed_policy_adapter_validates_external_output(): + context = SpeedContext(0.0, 0, 100, (1.0, 1.5)) + policy = SpeedPolicyAdapter(lambda observation, metadata: 1.5) + assert policy(np.zeros(3), context) == 1.5 + + invalid = SpeedPolicyAdapter(lambda observation, metadata: 0.0) + with pytest.raises(ValueError, match="positive"): + invalid(np.zeros(3), context) + + +def test_profile_policy_segments_nominal_time(): + policy = SpeedProfilePolicy([1.0, 2.0, 3.0]) + assert policy.select_speed(None, SpeedContext(0, 0, 90, (1.0,))) == 1.0 + assert policy.select_speed(None, SpeedContext(40, 0, 90, (1.0,))) == 2.0 + assert policy.select_speed(None, SpeedContext(89, 0, 90, (1.0,))) == 3.0 + + +def test_fixed_speed_policy_pairs_with_scripted_environment(): + env = create_speed_env("tea_bag", seed=0) + result = rollout_speed_policy(env, FixedSpeedPolicy(1.5)) + assert result["success"] + assert result["physics_steps"] < env.episode_len + + +def test_physics_instability_becomes_failed_terminal_transition(monkeypatch): + env = create_speed_env("tea_bag", seed=0) + env.reset() + + def unstable_step(action): + del action + raise control.PhysicsError("invalid simulated state") + + monkeypatch.setattr(env.env, "step", unstable_step) + observation, reward, done, info = env.step(3.0, quantized=False) + + assert done + assert not info["success"] + assert "invalid simulated state" in info["physics_error"] + assert observation.shape == (env.obs_space,) + assert np.isfinite(reward) + + +def test_recorded_chunk_policy_pairs_with_speed_environment(): + env = create_recorded_chunk_speed_env("tea_bag", chunk_size=25, seed=0) + result = rollout_speed_policy(env, FixedSpeedPolicy(1.0)) + assert result["success"] + assert result["mean_speed"] == 1.0 + + +@pytest.mark.rl +def test_prioritized_replay_samples_newest_transition(): + pytest.importorskip("torch") + from rl.rainbowDQN.replayBuffer import PrioritizedReplayBuffer + + replay = PrioritizedReplayBuffer( + obs_dim=1, + size=8, + batch_size=4, + alpha=1.0, + ) + for value in range(4): + replay.store( + np.array([value], dtype=np.float32), + value, + 0.0, + np.array([value + 1], dtype=np.float32), + False, + ) + replay.update_priorities( + np.arange(4), + np.array([1e-6, 1e-6, 1e-6, 100.0]), + ) + + assert 3 in replay.sample_batch(beta=0.4)["indices"] + + +@pytest.mark.rl +def test_public_rainbow_training_checkpoint_round_trip(tmp_path: Path): + pytest.importorskip("torch") + from speed_policy import RainbowSpeedPolicy + from speed_training import ( + RainbowTrainingConfig, + evaluate_rainbow_speed_policy, + train_rainbow_speed_policy, + ) + + env = create_speed_env("tea_bag", speed_values=(1.0,), seed=0) + config = RainbowTrainingConfig( + decisions=12, + memory_size=64, + batch_size=4, + learning_starts=4, + frame_skip=50, + gradient_steps=1, + train_interval=1, + target_update=2, + norm_update_interval=2, + exploration_steps=20, + atom_size=11, + n_step=3, + hidden_dim=32, + update_schedule="episode", + checkpoint_interval=6, + ) + checkpoint = tmp_path / "speed.pt" + result = train_rainbow_speed_policy( + env, + checkpoint, + config=config, + seed=0, + device="cpu", + progress=False, + ) + assert checkpoint.exists() + assert len(result["numbered_checkpoints"]) == 2 + assert all(Path(path).exists() for path in result["numbered_checkpoints"]) + assert result["updates"] > 0 + assert result["losses_finite"] + + observation = env.reset() + policy = RainbowSpeedPolicy.load(checkpoint) + assert policy.frame_skip == 50 + assert policy.observation_spec == env.observation_spec() + assert policy.environment_spec == env.environment_spec() + assert policy.select_speed( + observation, + SpeedContext(0, 0, env.episode_len, env.speed_values), + ) in env.speed_values + evaluation = evaluate_rainbow_speed_policy( + env, checkpoint, episodes=1, device="cpu" + ) + assert evaluation["episodes"] == 1 + assert "mean_acceleration" in evaluation + + randomized_env = create_speed_env( + "tea_bag", + speed_values=(1.0,), + randomize_object_pose=True, + seed=0, + ) + with pytest.raises(ValueError, match="Checkpoint environment"): + rollout_speed_policy(randomized_env, policy) diff --git a/uv.lock b/uv.lock new file mode 100644 index 0000000..bd18e50 --- /dev/null +++ b/uv.lock @@ -0,0 +1,911 @@ +version = 1 +revision = 3 +requires-python = "==3.10.*" + +[[package]] +name = "absl-py" +version = "2.5.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d0/4f/d79676ab82f2e42fc3611618139f13a9c4c31d0cff4b486982047679a802/absl_py-2.5.0.tar.gz", hash = "sha256:0c996f25c0490700fadabe6351630f6111534fa0ae252cc6d2014ea3b141135f", size = 118119, upload-time = "2026-07-03T10:57:48.157Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/58/0a/a10b45aab35b175aded078a462dc8d0c698f5b13946e7cb0869097b78bb6/absl_py-2.5.0-py3-none-any.whl", hash = "sha256:0f17b89f2a4eaaedc4f28c622998aa690564b3012a396a4ffad0821007fe03ba", size = 137410, upload-time = "2026-07-03T10:57:46.735Z" }, +] + +[[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 = "certifi" +version = "2026.7.22" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a3/c2/24167ea9858356b47a87a50d39908bfdb72ceeefe0041586e704e5376b3a/certifi-2026.7.22.tar.gz", hash = "sha256:741e2c3b351ddf169a738da9f2c048608ff7f2c5cc02f1ebc6b118bb090d5d55", size = 138112, upload-time = "2026-07-22T03:35:12.644Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0b/a7/71ac2cff56fec219ed242bb11b8efb69fcc4bec75db06fb7bfe35de520e6/certifi-2026.7.22-py3-none-any.whl", hash = "sha256:62f22742b58a1a33014a2b6b706588a8d7e2a88ae7bd1a6ebe8c992928483775", size = 136983, upload-time = "2026-07-22T03:35:11.276Z" }, +] + +[[package]] +name = "charset-normalizer" +version = "3.4.9" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/bd/2a/23f34ec9d04624958e137efdc394888716353190e75f25dd22c7a2c7a8aa/charset_normalizer-3.4.9.tar.gz", hash = "sha256:673611bbd43f0810bec0b0f028ddeaaa501190339cac411f347ac76917c3ae7b", size = 152439, upload-time = "2026-07-07T14:34:58.454Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ad/81/8e983840c6e5b93b33c2ba81aa3d52c2e42f0e9a690ce7607a2e61da4a5c/charset_normalizer-3.4.9-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:cd6280cf040f233bd7d3407b743b4b4c74f70e8e1c4199cb112a62c941c0772a", size = 322240, upload-time = "2026-07-07T14:32:36.236Z" }, + { url = "https://files.pythonhosted.org/packages/de/d1/b4319dc3229d8272fba305e206fc0a148e2de8d4087917ce62ae6382f359/charset_normalizer-3.4.9-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:aa99adc8f081b475a12843953db36831eaf83ec33eb46a90629ca6a5de45a616", size = 216475, upload-time = "2026-07-07T14:32:38.142Z" }, + { url = "https://files.pythonhosted.org/packages/80/33/6c99c1b3e6b8bf730e1bc809b9a2608f224145069114c479a2e9e1494346/charset_normalizer-3.4.9-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c1225416b463483160e4af85d5fc3a9690ccb53fd4b1865a6437825f5ede3209", size = 238670, upload-time = "2026-07-07T14:32:39.658Z" }, + { url = "https://files.pythonhosted.org/packages/7f/f4/ffbb83546e1f198ecc70ecd372b65cf2b50f9068b380abd67640f17a8e18/charset_normalizer-3.4.9-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:16d10d789dd9bcca1173c95af82c58433122564b7bc39385124be735a35cbe99", size = 233476, upload-time = "2026-07-07T14:32:41.155Z" }, + { url = "https://files.pythonhosted.org/packages/e8/5f/b98b8da398637b551e427e7be922bdec19177dc54d6811dcdaa503f23aac/charset_normalizer-3.4.9-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9bb41182d93ea91f60b4bc8fbf4c820c69ef8a12ab2d917f3f1834f1acad07e8", size = 223817, upload-time = "2026-07-07T14:32:42.592Z" }, + { url = "https://files.pythonhosted.org/packages/36/31/a276bb2e66243072a3fd06fdcab9cbb61a305b02143d70d2bda21d888fa8/charset_normalizer-3.4.9-cp310-cp310-manylinux_2_31_armv7l.whl", hash = "sha256:bcf74c1df76758a395bf0af608c04c82257523f55c9868b334f06270d0f2112b", size = 207974, upload-time = "2026-07-07T14:32:44.258Z" }, + { url = "https://files.pythonhosted.org/packages/5e/be/7ee4453d7e88dfbc4104ccd34900b9f2c7c17dac22881865fe0e82424a25/charset_normalizer-3.4.9-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b5314963fce9b0b12743891de876e724997864ee22aa496f903f426c7e2fa5b2", size = 221655, upload-time = "2026-07-07T14:32:45.64Z" }, + { url = "https://files.pythonhosted.org/packages/1d/85/181c652953eb5276d198f375b1dd641047392050098100a3a02d6534f657/charset_normalizer-3.4.9-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:e9701d0049d92c16703a42771b98d560b95248949f23f8cf7b4eddd201814fb9", size = 219229, upload-time = "2026-07-07T14:32:47.376Z" }, + { url = "https://files.pythonhosted.org/packages/0c/e7/aaf6da33fc9f4691cda8f7efbc9f69179d3d39ec8a4799baf273ee1d8db0/charset_normalizer-3.4.9-cp310-cp310-musllinux_1_2_armv7l.whl", hash = "sha256:65a7ff3f705e57d392f7261b6d0550fe137c3019477431f1c355e0db0a7d3e15", size = 209704, upload-time = "2026-07-07T14:32:48.855Z" }, + { url = "https://files.pythonhosted.org/packages/63/01/f2fb3bd3a73be48b173ee0c6aa8d2497af97d5663a8c4c4b491de4c62f7a/charset_normalizer-3.4.9-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:79580094b00d1789d1f93ea55bc43cb2f611910c72235b7657f3482ddcc1b22d", size = 226243, upload-time = "2026-07-07T14:32:50.239Z" }, + { url = "https://files.pythonhosted.org/packages/c4/02/c57a22739fe05246b0b5783b3bfb6afaac4eebb46f3ececdfb2f048f780e/charset_normalizer-3.4.9-cp310-cp310-win32.whl", hash = "sha256:432786d3561e69aeeae6c7e8648964ce0ad05736120135601f87ac26b9c83381", size = 150935, upload-time = "2026-07-07T14:32:51.676Z" }, + { url = "https://files.pythonhosted.org/packages/37/8d/ca39a7559a4797505530d084fd3a49a2c959efbbbff146302fb7be4e3b35/charset_normalizer-3.4.9-cp310-cp310-win_amd64.whl", hash = "sha256:8c041122946b7ba21bb32c45b1aa57b1be35527690aeb3c5c234521085632eee", size = 162314, upload-time = "2026-07-07T14:32:53.193Z" }, + { url = "https://files.pythonhosted.org/packages/01/da/a44bd7a13d426e69e4894557106cd58669097bfad4a8681123b618fbfc5d/charset_normalizer-3.4.9-cp310-cp310-win_arm64.whl", hash = "sha256:375b83ed0aecfce76c16d198fbc21f3b11b337d68662bea0a995046682a11419", size = 153075, upload-time = "2026-07-07T14:32:54.554Z" }, + { url = "https://files.pythonhosted.org/packages/98/2b/f97f1c193fb855c345d678f5077d6926034db0722df74c8f057020e05a25/charset_normalizer-3.4.9-py3-none-any.whl", hash = "sha256:68e5f26a1ad57ded6d1cfb85331d1c1a195314756471d97758c48498bb4dcdf5", size = 64538, upload-time = "2026-07-07T14:34:56.993Z" }, +] + +[[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 = "contourpy" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/54/eb9bfc647b19f2009dd5c7f5ec51c4e6ca831725f1aea7a993034f483147/contourpy-1.3.2.tar.gz", hash = "sha256:b6945942715a034c671b7fc54f9588126b0b8bf23db2696e3ca8328f3ff0ab54", size = 13466130, upload-time = "2025-04-15T17:47:53.79Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/12/a3/da4153ec8fe25d263aa48c1a4cbde7f49b59af86f0b6f7862788c60da737/contourpy-1.3.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:ba38e3f9f330af820c4b27ceb4b9c7feee5fe0493ea53a8720f4792667465934", size = 268551, upload-time = "2025-04-15T17:34:46.581Z" }, + { url = "https://files.pythonhosted.org/packages/2f/6c/330de89ae1087eb622bfca0177d32a7ece50c3ef07b28002de4757d9d875/contourpy-1.3.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:dc41ba0714aa2968d1f8674ec97504a8f7e334f48eeacebcaa6256213acb0989", size = 253399, upload-time = "2025-04-15T17:34:51.427Z" }, + { url = "https://files.pythonhosted.org/packages/c1/bd/20c6726b1b7f81a8bee5271bed5c165f0a8e1f572578a9d27e2ccb763cb2/contourpy-1.3.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9be002b31c558d1ddf1b9b415b162c603405414bacd6932d031c5b5a8b757f0d", size = 312061, upload-time = "2025-04-15T17:34:55.961Z" }, + { url = "https://files.pythonhosted.org/packages/22/fc/a9665c88f8a2473f823cf1ec601de9e5375050f1958cbb356cdf06ef1ab6/contourpy-1.3.2-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8d2e74acbcba3bfdb6d9d8384cdc4f9260cae86ed9beee8bd5f54fee49a430b9", size = 351956, upload-time = "2025-04-15T17:35:00.992Z" }, + { url = "https://files.pythonhosted.org/packages/25/eb/9f0a0238f305ad8fb7ef42481020d6e20cf15e46be99a1fcf939546a177e/contourpy-1.3.2-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e259bced5549ac64410162adc973c5e2fb77f04df4a439d00b478e57a0e65512", size = 320872, upload-time = "2025-04-15T17:35:06.177Z" }, + { url = "https://files.pythonhosted.org/packages/32/5c/1ee32d1c7956923202f00cf8d2a14a62ed7517bdc0ee1e55301227fc273c/contourpy-1.3.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ad687a04bc802cbe8b9c399c07162a3c35e227e2daccf1668eb1f278cb698631", size = 325027, upload-time = "2025-04-15T17:35:11.244Z" }, + { url = "https://files.pythonhosted.org/packages/83/bf/9baed89785ba743ef329c2b07fd0611d12bfecbedbdd3eeecf929d8d3b52/contourpy-1.3.2-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:cdd22595308f53ef2f891040ab2b93d79192513ffccbd7fe19be7aa773a5e09f", size = 1306641, upload-time = "2025-04-15T17:35:26.701Z" }, + { url = "https://files.pythonhosted.org/packages/d4/cc/74e5e83d1e35de2d28bd97033426b450bc4fd96e092a1f7a63dc7369b55d/contourpy-1.3.2-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:b4f54d6a2defe9f257327b0f243612dd051cc43825587520b1bf74a31e2f6ef2", size = 1374075, upload-time = "2025-04-15T17:35:43.204Z" }, + { url = "https://files.pythonhosted.org/packages/0c/42/17f3b798fd5e033b46a16f8d9fcb39f1aba051307f5ebf441bad1ecf78f8/contourpy-1.3.2-cp310-cp310-win32.whl", hash = "sha256:f939a054192ddc596e031e50bb13b657ce318cf13d264f095ce9db7dc6ae81c0", size = 177534, upload-time = "2025-04-15T17:35:46.554Z" }, + { url = "https://files.pythonhosted.org/packages/54/ec/5162b8582f2c994721018d0c9ece9dc6ff769d298a8ac6b6a652c307e7df/contourpy-1.3.2-cp310-cp310-win_amd64.whl", hash = "sha256:c440093bbc8fc21c637c03bafcbef95ccd963bc6e0514ad887932c18ca2a759a", size = 221188, upload-time = "2025-04-15T17:35:50.064Z" }, + { url = "https://files.pythonhosted.org/packages/33/05/b26e3c6ecc05f349ee0013f0bb850a761016d89cec528a98193a48c34033/contourpy-1.3.2-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:fd93cc7f3139b6dd7aab2f26a90dde0aa9fc264dbf70f6740d498a70b860b82c", size = 265681, upload-time = "2025-04-15T17:44:59.314Z" }, + { url = "https://files.pythonhosted.org/packages/2b/25/ac07d6ad12affa7d1ffed11b77417d0a6308170f44ff20fa1d5aa6333f03/contourpy-1.3.2-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:107ba8a6a7eec58bb475329e6d3b95deba9440667c4d62b9b6063942b61d7f16", size = 315101, upload-time = "2025-04-15T17:45:04.165Z" }, + { url = "https://files.pythonhosted.org/packages/8f/4d/5bb3192bbe9d3f27e3061a6a8e7733c9120e203cb8515767d30973f71030/contourpy-1.3.2-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:ded1706ed0c1049224531b81128efbd5084598f18d8a2d9efae833edbd2b40ad", size = 220599, upload-time = "2025-04-15T17:45:08.456Z" }, +] + +[[package]] +name = "cycler" +version = "0.12.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a9/95/a3dbbb5028f35eafb79008e7522a75244477d2838f38cbb722248dabc2a8/cycler-0.12.1.tar.gz", hash = "sha256:88bb128f02ba341da8ef447245a9e138fae777f6a23943da4540077d3601eb1c", size = 7615, upload-time = "2023-10-07T05:32:18.335Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, +] + +[[package]] +name = "dm-control" +version = "1.0.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "dm-env" }, + { name = "dm-tree" }, + { name = "glfw" }, + { name = "labmaze" }, + { name = "lxml" }, + { name = "mujoco" }, + { name = "numpy" }, + { name = "protobuf" }, + { name = "pyopengl" }, + { name = "pyparsing" }, + { name = "requests" }, + { name = "scipy" }, + { name = "setuptools" }, + { name = "tqdm" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/46/bb/48499d7a063652250b5af00b3193e99b43fb24586d45b2e3dcaa61bcfb21/dm_control-1.0.9.tar.gz", hash = "sha256:452f175bd04fdd0692c01cfa25dc09ff7ce61a81674f92cb1c4ac0f5cb6e1fb0", size = 38980591, upload-time = "2022-12-09T13:20:05.865Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/16/34/e9d3224a0c4496e62cd6b6eb53fe863644ebfff19bad1c8e4fb0d26270b5/dm_control-1.0.9-py3-none-any.whl", hash = "sha256:531e0bb310c62fd9e4058a295fae4f3f22a58de3e7c0e394fe08bbd8db2ba331", size = 39281788, upload-time = "2022-12-09T13:19:52.584Z" }, +] + +[[package]] +name = "dm-env" +version = "1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "dm-tree" }, + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/62/c9/93e8d6239d5806508a2ee4b370e67c6069943ca149f59f533923737a99b7/dm-env-1.6.tar.gz", hash = "sha256:a436eb1c654c39e0c986a516cee218bea7140b510fceff63f97eb4fcff3d93de", size = 20187, upload-time = "2022-12-21T00:25:29.306Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/08/7e/36d548040e61337bf9182637a589c44da407a47a923ee88aec7f0e89867c/dm_env-1.6-py3-none-any.whl", hash = "sha256:0eabb6759dd453b625e041032f7ae0c1e87d4eb61b6a96b9ca586483837abf29", size = 26339, upload-time = "2022-12-21T00:25:37.128Z" }, +] + +[[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/52/76/781bc1a8ce0f4be153755e36be547d7d36964fb6d265b9b29503ed3e0a0f/dm_tree-0.1.10-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:b8b606661abcb5336e60ce4cd6f3bec794ec794f91d8de001c1ea451dd7a7411", size = 311543, upload-time = "2026-03-31T17:35:04.766Z" }, + { url = "https://files.pythonhosted.org/packages/7c/b9/2f6278c07728c60411d363aab1b5de53b75bb3bdf27032e871c240f3b4de/dm_tree-0.1.10-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0fcf11edc379723b8a17b76b6f63643e9280d55f23c37994805bf27282074192", size = 181202, upload-time = "2026-03-31T17:35:06.266Z" }, + { url = "https://files.pythonhosted.org/packages/6d/d6/2b88e4eb5e3e5ecf9d61afe55e463ce7c9e4d1dc6630b9f7038a62e548d1/dm_tree-0.1.10-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7644b24bb5c7810b601f4b28cf6e4b516da96713ebfd9e45585240318c6b9384", size = 183874, upload-time = "2026-03-31T17:35:07.746Z" }, + { url = "https://files.pythonhosted.org/packages/49/9f/6fea8ba8cb69136af0511868bc41dfb88d0b8c3a85497dbbf16271cd84a7/dm_tree-0.1.10-cp310-cp310-win_amd64.whl", hash = "sha256:d7f42b2148a1a3758230fbf93c06b9475e5d4bd62a4d19eacafa8210b03421ac", size = 110740, upload-time = "2026-03-31T17:35:09.045Z" }, +] + +[[package]] +name = "einops" +version = "0.8.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2c/77/850bef8d72ffb9219f0b1aac23fbc1bf7d038ee6ea666f331fa273031aa2/einops-0.8.2.tar.gz", hash = "sha256:609da665570e5e265e27283aab09e7f279ade90c4f01bcfca111f3d3e13f2827", size = 56261, upload-time = "2026-01-26T04:13:17.638Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/09/f8d8f8f31e4483c10a906437b4ce31bdf3d6d417b73fe33f1a8b59e34228/einops-0.8.2-py3-none-any.whl", hash = "sha256:54058201ac7087911181bfec4af6091bb59380360f069276601256a76af08193", size = 65638, upload-time = "2026-01-26T04:13:18.546Z" }, +] + +[[package]] +name = "exceptiongroup" +version = "1.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8a/0e/97c33bf5009bdbac74fd2beace167cab3f978feb69cc36f1ef79360d6c4e/exceptiongroup-1.3.1-py3-none-any.whl", hash = "sha256:a7a39a3bd276781e98394987d3a5701d0c4edffb633bb7a5144577f82c773598", size = 16740, upload-time = "2025-11-21T23:01:53.443Z" }, +] + +[[package]] +name = "filelock" +version = "3.32.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/57/3ba6e6cb097f85b855b00163d169f35365f44277df044dcf96d55b8f62a3/filelock-3.32.2.tar.gz", hash = "sha256:c33351e1f49cae33414acbc6d56784e6ecee82514ec90795da1161fc4836b5b8", size = 217172, upload-time = "2026-07-29T22:46:04.895Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/e8/72f8cef9fdfeffe06213fe8508039396ee48daa0e3259457ed766173bfd6/filelock-3.32.2-py3-none-any.whl", hash = "sha256:87dd94cf281e586d135fa51132b8e3d9a598b316e90377a288663c9321036c82", size = 98830, upload-time = "2026-07-29T22:46:03.52Z" }, +] + +[[package]] +name = "fonttools" +version = "4.63.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/84/69/c97f2c18e0db87d2c7b15da1974dace76ae938f1cfa22e2727a648b7ed43/fonttools-4.63.0.tar.gz", hash = "sha256:caeb583deeb5168e694b65cda8b4ee62abedfa66cf88488734466f2366b9c4e0", size = 3597189, upload-time = "2026-05-14T12:04:30.958Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f2/c9/4141c90a90db20f807c7e10bfd689fe53eb8f7f4caff58ee4d4dfe46919f/fonttools-4.63.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:e3297a6a4059b4acc3a1e9a8b04741f240a80044eef08ebd32e8b5bcdddce75b", size = 2884632, upload-time = "2026-05-14T12:02:38.56Z" }, + { url = "https://files.pythonhosted.org/packages/b8/46/ad12b5c10eae602d7ef814b02afa08aacbf89da917fed5b071282b7eadc2/fonttools-4.63.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:b1cd75a03ad8cb5bc40c90bfde68c0c47de423aa19e5c0f362b43520645eea94", size = 2429441, upload-time = "2026-05-14T12:02:41.162Z" }, + { url = "https://files.pythonhosted.org/packages/90/8f/bdca24a84c81d56fffed052229cdcff368f6e05882e526f4558891481f65/fonttools-4.63.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c0425b277a59cff3d80ca42162a8de360f318438a2ac83570842a678d826d579", size = 4946346, upload-time = "2026-05-14T12:02:43.41Z" }, + { url = "https://files.pythonhosted.org/packages/04/59/a639c0e136441ee91a65b56fdf89e5d075927e7a09c559d1b0f5276577db/fonttools-4.63.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d7e5c9973aa04c95650c96e5f5ad865fbf42d62079163ecfab1e01cbc2504c22", size = 4903184, upload-time = "2026-05-14T12:02:45.742Z" }, + { url = "https://files.pythonhosted.org/packages/e6/53/91b7e0cb45b536f3da1b29ba8cbab89f27e8b986809e0b1982303a3f4eca/fonttools-4.63.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:cb014d58140a38135f16064c74c652ed57aa0b75cbf8bb59cac821f7edb5334e", size = 4922967, upload-time = "2026-05-14T12:02:48.386Z" }, + { url = "https://files.pythonhosted.org/packages/c7/b7/87439bf44e6b97c5538cd29d0b7e366a5b8ce2cc132a4134fb67fa3f2fa2/fonttools-4.63.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:032038247a96c1690f9f31e377c389383c902531b085aa4e4dabd6f57f870e69", size = 5042799, upload-time = "2026-05-14T12:02:50.424Z" }, + { url = "https://files.pythonhosted.org/packages/ad/7c/8b96c3263b89ef99cded544c0f0636686f85dbd3c211c4dceef0231fca23/fonttools-4.63.0-cp310-cp310-win32.whl", hash = "sha256:a8b33a82979e0a6a34ff435cc81317be1f95ec1ebb7a3a2d1c8a6a54f02ae44e", size = 1519704, upload-time = "2026-05-14T12:02:52.523Z" }, + { url = "https://files.pythonhosted.org/packages/e5/4d/2c2f0069970b6907de8fb5b05c5c0193cc22f717df151d1c7aef1c738f58/fonttools-4.63.0-cp310-cp310-win_amd64.whl", hash = "sha256:0c18358a155d75034911c5ee397a5b44cd19dd325dbb8b35fb60bf421d6a72ac", size = 1568666, upload-time = "2026-05-14T12:02:54.917Z" }, + { url = "https://files.pythonhosted.org/packages/2c/47/c99d5268f354002ce80f8d029cd9d7d872969da1de8b93d32de4dc56d6f4/fonttools-4.63.0-py3-none-any.whl", hash = "sha256:445af2eab030a16b9171ea8bdda7ebf7d96bda2df88ee182a464252f6e05e20d", size = 1164562, upload-time = "2026-05-14T12:04:29.092Z" }, +] + +[[package]] +name = "fsspec" +version = "2026.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/00/78/f34251dadb8f3921264a1d9b8946f5e542014ee2614b285261b4e40e6775/fsspec-2026.7.0.tar.gz", hash = "sha256:c803c40f4cf860b49dea58ee3e1c33cb9c790520e233537e1340049f89b82a88", size = 317040, upload-time = "2026-07-28T16:34:51.052Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl", hash = "sha256:b57ddbafedfaef7018c1ecab32aa200a9d7ca26b77965f64e48b70061249d279", size = 206583, upload-time = "2026-07-28T16:34:49.538Z" }, +] + +[[package]] +name = "glfw" +version = "2.10.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/62/096058bcb4b4fb28f7ecd28fb048f07969d90b243c417af5f6d09d45a0c2/glfw-2.10.2.tar.gz", hash = "sha256:5d2cf97c66bc42a6b583be0e307eae5a3945438322e2ed0c5e4f14dc251d693d", size = 36307, upload-time = "2026-07-21T14:42:36.732Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c4/80/5e575ec14f54dc0a644e7215985b9ee7c5044fcc3594d6b83dfdeb04ccd4/glfw-2.10.2-py2.py3-none-macosx_10_6_intel.whl", hash = "sha256:a04d25bb535a3a29e173ea1abd658e161681b5d3816e00bc51c122f74cd7fa3a", size = 110295, upload-time = "2026-07-21T14:42:25.625Z" }, + { url = "https://files.pythonhosted.org/packages/e5/2e/22216db429690aa51fa87d2e2b5e6b610c5a37b9df7768da2927f52f8704/glfw-2.10.2-py2.py3-none-macosx_11_0_arm64.whl", hash = "sha256:c000b8a5e5fb4374b2c18d12347255ccb92001b9bcf80ef74c56072acbb98fe3", size = 107146, upload-time = "2026-07-21T14:42:26.854Z" }, + { url = "https://files.pythonhosted.org/packages/75/a4/9ccee9d5c3b9c5d6bafce3e8d53ab8318eff30124e3c19b95580bebaa6d4/glfw-2.10.2-py2.py3-none-manylinux2014_aarch64.whl", hash = "sha256:0ffe2b48b51f37b05e8f027f57c0b8cd213affaf22fc83401bd9621692d4c6c8", size = 235005, upload-time = "2026-07-21T14:42:27.955Z" }, + { url = "https://files.pythonhosted.org/packages/77/44/37aeed50c581c76e1c5bb83b696a06f98f8ecc979ed8564dde3eef908e8b/glfw-2.10.2-py2.py3-none-manylinux2014_x86_64.whl", hash = "sha256:82daf1f87b2a48815637dc6f6720d4a55a57ba1d3683322dc20dc28065754f97", size = 246951, upload-time = "2026-07-21T14:42:29.107Z" }, + { url = "https://files.pythonhosted.org/packages/52/fc/29b4f1a8c8e33d8934781330a0845458b01da7f389d9fe6a3350e752b7f3/glfw-2.10.2-py2.py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:7604425ec95665cc6a25c3ed7d07e9660d394a34472ceffaec0843c844e7b649", size = 236018, upload-time = "2026-07-21T14:42:30.623Z" }, + { url = "https://files.pythonhosted.org/packages/69/84/1b98671314ea5f40397586b6cd14db913de3196333be558fdc177e61708f/glfw-2.10.2-py2.py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:b860de3ca0686182483f98f3ddd12e660acf25b3e0d521450ec9a3f999f72a65", size = 248490, upload-time = "2026-07-21T14:42:32.377Z" }, + { url = "https://files.pythonhosted.org/packages/12/d9/bd2e6c8dbadcf69fcd2289705e75881258b46221c196903a52e9f6e97322/glfw-2.10.2-py2.py3-none-win32.whl", hash = "sha256:a94aefe8c48886fd83cd6e11e936846166a4b5b94abf1f4bf599d984774b0b1b", size = 557654, upload-time = "2026-07-21T14:42:33.903Z" }, + { url = "https://files.pythonhosted.org/packages/e0/6c/4ca5f3ab85a8d7f612ce857208726e9b834e54e559751e3d4f98cc2f7509/glfw-2.10.2-py2.py3-none-win_amd64.whl", hash = "sha256:15fd0666cd8f1b0ecb535c3abb712b8ddda053bf8d16de2353b1c6ced0e65402", size = 564433, upload-time = "2026-07-21T14:42:35.465Z" }, +] + +[[package]] +name = "idna" +version = "3.18" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cd/63/9496c57188a2ee585e0f1db071d75089a11e98aa86eb99d9d7618fc1edce/idna-3.18.tar.gz", hash = "sha256:ffb385a7e039654cef1ab9ef32c6fafe283c0c0467bba1d9029738ce4a14a848", size = 196711, upload-time = "2026-06-02T14:34:07.794Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/5e/d4e9f1a599fb8e573b7b87160658329fbf28d19eac2718f51fc3def3aa5a/idna-3.18-py3-none-any.whl", hash = "sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2", size = 65455, upload-time = "2026-06-02T14:34:06.319Z" }, +] + +[[package]] +name = "imageio" +version = "2.37.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "pillow" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/48/62/aa770a9307508d2a2a2c62d536a49347bffe9e55322db27838d3c93d0b07/imageio-2.37.4.tar.gz", hash = "sha256:e45cbc5e83502047fb138f7f585f7f105a136a57eea5f4b3cfc6ce1b52720bd3", size = 390173, upload-time = "2026-07-20T05:26:11.369Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3e/2d/ca050652104bab2cf55e569db2a178b1b61cb041fef28307f2db383f6d9f/imageio-2.37.4-py3-none-any.whl", hash = "sha256:1ab2e22c8debf700f24c3ac43e8f95f3b3a8110c83b93411e97b4b0b2cd1c7e6", size = 318000, upload-time = "2026-07-20T05:26:09.874Z" }, +] + +[[package]] +name = "imageio-ffmpeg" +version = "0.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/44/bd/c3343c721f2a1b0c9fc71c1aebf1966a3b7f08c2eea8ed5437a2865611d6/imageio_ffmpeg-0.6.0.tar.gz", hash = "sha256:e2556bed8e005564a9f925bb7afa4002d82770d6b08825078b7697ab88ba1755", size = 25210, upload-time = "2025-01-16T21:34:32.747Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/58/87ef68ac83f4c7690961bce288fd8e382bc5f1513860fc7f90a9c1c1c6bf/imageio_ffmpeg-0.6.0-py3-none-macosx_10_9_intel.macosx_10_9_x86_64.whl", hash = "sha256:9d2baaf867088508d4a3458e61eeb30e945c4ad8016025545f66c4b5aaef0a61", size = 24932969, upload-time = "2025-01-16T21:34:20.464Z" }, + { url = "https://files.pythonhosted.org/packages/40/5c/f3d8a657d362cc93b81aab8feda487317da5b5d31c0e1fdfd5e986e55d17/imageio_ffmpeg-0.6.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:b1ae3173414b5fc5f538a726c4e48ea97edc0d2cdc11f103afee655c463fa742", size = 21113891, upload-time = "2025-01-16T21:34:00.277Z" }, + { url = "https://files.pythonhosted.org/packages/33/e7/1925bfbc563c39c1d2e82501d8372734a5c725e53ac3b31b4c2d081e895b/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:1d47bebd83d2c5fc770720d211855f208af8a596c82d17730aa51e815cdee6dc", size = 25632706, upload-time = "2025-01-16T21:33:53.475Z" }, + { url = "https://files.pythonhosted.org/packages/a0/2d/43c8522a2038e9d0e7dbdf3a61195ecc31ca576fb1527a528c877e87d973/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:c7e46fcec401dd990405049d2e2f475e2b397779df2519b544b8aab515195282", size = 29498237, upload-time = "2025-01-16T21:34:13.726Z" }, + { url = "https://files.pythonhosted.org/packages/a0/13/59da54728351883c3c1d9fca1710ab8eee82c7beba585df8f25ca925f08f/imageio_ffmpeg-0.6.0-py3-none-win32.whl", hash = "sha256:196faa79366b4a82f95c0f4053191d2013f4714a715780f0ad2a68ff37483cc2", size = 19652251, upload-time = "2025-01-16T21:34:06.812Z" }, + { url = "https://files.pythonhosted.org/packages/2c/c6/fa760e12a2483469e2bf5058c5faff664acf66cadb4df2ad6205b016a73d/imageio_ffmpeg-0.6.0-py3-none-win_amd64.whl", hash = "sha256:02fa47c83703c37df6bfe4896aab339013f62bf02c5ebf2dce6da56af04ffc0a", size = 31246824, upload-time = "2025-01-16T21:34:28.6Z" }, +] + +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, +] + +[[package]] +name = "kiwisolver" +version = "1.5.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d0/67/9c61eccb13f0bdca9307614e782fec49ffdde0f7a2314935d489fa93cd9c/kiwisolver-1.5.0.tar.gz", hash = "sha256:d4193f3d9dc3f6f79aaed0e5637f45d98850ebf01f7ca20e69457f3e8946b66a", size = 103482, upload-time = "2026-03-09T13:15:53.382Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ac/f8/06549565caa026e540b7e7bab5c5a90eb7ca986015f4c48dace243cd24d9/kiwisolver-1.5.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:32cc0a5365239a6ea0c6ed461e8838d053b57e397443c0ca894dcc8e388d4374", size = 122802, upload-time = "2026-03-09T13:12:37.515Z" }, + { url = "https://files.pythonhosted.org/packages/84/eb/8476a0818850c563ff343ea7c9c05dcdcbd689a38e01aa31657df01f91fa/kiwisolver-1.5.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:cc0b66c1eec9021353a4b4483afb12dfd50e3669ffbb9152d6842eb34c7e29fd", size = 66216, upload-time = "2026-03-09T13:12:38.812Z" }, + { url = "https://files.pythonhosted.org/packages/f3/c4/f9c8a6b4c21aed4198566e45923512986d6cef530e7263b3a5f823546561/kiwisolver-1.5.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:86e0287879f75621ae85197b0877ed2f8b7aa57b511c7331dce2eb6f4de7d476", size = 63917, upload-time = "2026-03-09T13:12:40.053Z" }, + { url = "https://files.pythonhosted.org/packages/f1/0e/ba4ae25d03722f64de8b2c13e80d82ab537a06b30fc7065183c6439357e3/kiwisolver-1.5.0-cp310-cp310-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:62f59da443c4f4849f73a51a193b1d9d258dcad0c41bc4d1b8fb2bcc04bfeb22", size = 1628776, upload-time = "2026-03-09T13:12:41.976Z" }, + { url = "https://files.pythonhosted.org/packages/8a/e4/3f43a011bc8a0860d1c96f84d32fa87439d3feedf66e672fef03bf5e8bac/kiwisolver-1.5.0-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9190426b7aa26c5229501fa297b8d0653cfd3f5a36f7990c264e157cbf886b3b", size = 1228164, upload-time = "2026-03-09T13:12:44.002Z" }, + { url = "https://files.pythonhosted.org/packages/4b/34/3a901559a1e0c218404f9a61a93be82d45cb8f44453ba43088644980f033/kiwisolver-1.5.0-cp310-cp310-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c8277104ded0a51e699c8c3aff63ce2c56d4ed5519a5f73e0fd7057f959a2b9e", size = 1246656, upload-time = "2026-03-09T13:12:45.557Z" }, + { url = "https://files.pythonhosted.org/packages/87/9e/f78c466ea20527822b95ad38f141f2de1dcd7f23fb8716b002b0d91bbe59/kiwisolver-1.5.0-cp310-cp310-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:8f9baf6f0a6e7571c45c8863010b45e837c3ee1c2c77fcd6ef423be91b21fedb", size = 1295562, upload-time = "2026-03-09T13:12:47.562Z" }, + { url = "https://files.pythonhosted.org/packages/0a/66/fd0e4a612e3a286c24e6d6f3a5428d11258ed1909bc530ba3b59807fd980/kiwisolver-1.5.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:cff8e5383db4989311f99e814feeb90c4723eb4edca425b9d5d9c3fefcdd9537", size = 2178473, upload-time = "2026-03-09T13:12:50.254Z" }, + { url = "https://files.pythonhosted.org/packages/dc/8e/6cac929e0049539e5ee25c1ee937556f379ba5204840d03008363ced662d/kiwisolver-1.5.0-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:ebae99ed6764f2b5771c522477b311be313e8841d2e0376db2b10922daebbba4", size = 2274035, upload-time = "2026-03-09T13:12:51.785Z" }, + { url = "https://files.pythonhosted.org/packages/ca/d3/9d0c18f1b52ea8074b792452cf17f1f5a56bd0302a85191f405cfbf9da16/kiwisolver-1.5.0-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:d5cd5189fc2b6a538b75ae45433140c4823463918f7b1617c31e68b085c0022c", size = 2443217, upload-time = "2026-03-09T13:12:53.329Z" }, + { url = "https://files.pythonhosted.org/packages/45/2a/6e19368803a038b2a90857bf4ee9e3c7b667216d045866bf22d3439fd75e/kiwisolver-1.5.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:f42c23db5d1521218a3276bb08666dcb662896a0be7347cba864eca45ff64ede", size = 2249196, upload-time = "2026-03-09T13:12:55.057Z" }, + { url = "https://files.pythonhosted.org/packages/75/2b/3f641dfcbe72e222175d626bacf2f72c3b34312afec949dd1c50afa400f5/kiwisolver-1.5.0-cp310-cp310-win_amd64.whl", hash = "sha256:94eff26096eb5395136634622515b234ecb6c9979824c1f5004c6e3c3c85ccd2", size = 73389, upload-time = "2026-03-09T13:12:56.496Z" }, + { url = "https://files.pythonhosted.org/packages/da/88/299b137b9e0025d8982e03d2d52c123b0a2b159e84b0ef1501ef446339cf/kiwisolver-1.5.0-cp310-cp310-win_arm64.whl", hash = "sha256:dd952e03bfbb096cfe2dd35cd9e00f269969b67536cb4370994afc20ff2d0875", size = 64782, upload-time = "2026-03-09T13:12:57.609Z" }, + { url = "https://files.pythonhosted.org/packages/17/6f/6fd4f690a40c2582fa34b97d2678f718acf3706b91d270c65ecb455d0a06/kiwisolver-1.5.0-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:295d9ffe712caa9f8a3081de8d32fc60191b4b51c76f02f951fd8407253528f4", size = 59606, upload-time = "2026-03-09T13:15:40.81Z" }, + { url = "https://files.pythonhosted.org/packages/82/a0/2355d5e3b338f13ce63f361abb181e3b6ea5fffdb73f739b3e80efa76159/kiwisolver-1.5.0-pp310-pypy310_pp73-macosx_11_0_arm64.whl", hash = "sha256:51e8c4084897de9f05898c2c2a39af6318044ae969d46ff7a34ed3f96274adca", size = 57537, upload-time = "2026-03-09T13:15:42.071Z" }, + { url = "https://files.pythonhosted.org/packages/c8/b9/1d50e610ecadebe205b71d6728fd224ce0e0ca6aba7b9cbe1da049203ac5/kiwisolver-1.5.0-pp310-pypy310_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b83af57bdddef03c01a9138034c6ff03181a3028d9a1003b301eb1a55e161a3f", size = 79888, upload-time = "2026-03-09T13:15:43.317Z" }, + { url = "https://files.pythonhosted.org/packages/cd/ee/b85ffcd75afed0357d74f0e6fc02a4507da441165de1ca4760b9f496390d/kiwisolver-1.5.0-pp310-pypy310_pp73-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bf4679a3d71012a7c2bf360e5cd878fbd5e4fcac0896b56393dec239d81529ed", size = 77584, upload-time = "2026-03-09T13:15:44.605Z" }, + { url = "https://files.pythonhosted.org/packages/6b/dd/644d0dde6010a8583b4cd66dd41c5f83f5325464d15c4f490b3340ab73b4/kiwisolver-1.5.0-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:41024ed50e44ab1a60d3fe0a9d15a4ccc9f5f2b1d814ff283c8d01134d5b81bc", size = 73390, upload-time = "2026-03-09T13:15:45.832Z" }, +] + +[[package]] +name = "labmaze" +version = "1.0.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "numpy" }, + { name = "setuptools" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/93/0a/139c4ae896b9413bd4ca69c62b08ee98dcfc78a9cbfdb7cadd0dce2ad31d/labmaze-1.0.6.tar.gz", hash = "sha256:2e8de7094042a77d6972f1965cf5c9e8f971f1b34d225752f343190a825ebe73", size = 4670455, upload-time = "2022-12-05T18:42:43.566Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/0c/6a3941f48644c0b9305c7a22bd51974be1fed8e9233b16c893d728805143/labmaze-1.0.6-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:b2ddef976dfd8d992b19cfa6c633f2eba7576d759c2082da534e3f727479a84a", size = 4815423, upload-time = "2022-12-05T18:41:47.351Z" }, + { url = "https://files.pythonhosted.org/packages/d0/fe/b038c6a15732eb064767dc92ca39a38b2f5df183576384f0cfb6a4840f69/labmaze-1.0.6-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:157efaa93228c8ccce5cae337902dd652093e0fba9d3a0f6506e4bee272bb66f", size = 4806825, upload-time = "2022-12-05T18:41:49.922Z" }, + { url = "https://files.pythonhosted.org/packages/59/ec/2762281d4f26845b20bb7529742a6914fcb07c8e7c522175b879df0127cf/labmaze-1.0.6-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b3ce98b9541c5fe6a306e411e7d018121dd646f2c9978d763fad86f9f30c5f57", size = 4871532, upload-time = "2022-12-05T18:41:52.784Z" }, + { url = "https://files.pythonhosted.org/packages/4d/93/abac7877e1d7de984a2f0f5be561ff0dc795ae7e22595cf2f7c7032cd27e/labmaze-1.0.6-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4e6433bd49bc541791de8191040526fddfebb77151620eb04203453f43ee486a", size = 4875892, upload-time = "2022-12-05T18:41:55.603Z" }, + { url = "https://files.pythonhosted.org/packages/c4/10/5262db11b3c1db8e4fbc3feed9baed4f95db6047b8d9dcaf4f9fb8da9ba3/labmaze-1.0.6-cp310-cp310-win_amd64.whl", hash = "sha256:6a507fc35961f1b1479708e2716f65e0d0611cefb55f31a77be29ce2339b6fef", size = 4812953, upload-time = "2022-12-05T18:41:58.098Z" }, +] + +[[package]] +name = "lxml" +version = "6.1.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/05/3b/aab6728cae887456f409b4d75e8a01856e4f04bd510de38052a47768b680/lxml-6.1.1.tar.gz", hash = "sha256:ba96ae44888e0185281e937633a743ea90d5a196c6000f82565ebb0580012d40", size = 4197430, upload-time = "2026-05-18T19:19:06.424Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/12/da/dbe4dfc01ac226fb0504fad035f4d69f3202f3502e20e68537631daddd96/lxml-6.1.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:09dd5b7075dc2f7709654a46543ba1ea3c2e217b2ed8fbd413a8a945a0f40f60", size = 8541124, upload-time = "2026-05-18T19:17:11.589Z" }, + { url = "https://files.pythonhosted.org/packages/78/20/f7095ed9fc2c025f9cfe71cc6ec9f1feb05624edc1812423b5f1aecf3d4b/lxml-6.1.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:f6ac4ef4d82dff54670227a69c67782ae0b811b5cf6b17954f1e8f7502fc0d1d", size = 4602783, upload-time = "2026-05-18T19:17:20.888Z" }, + { url = "https://files.pythonhosted.org/packages/4a/a4/65c63ca98bd129f6cff7b8c2fa48953ab058cc6005b541354e7dd54d8000/lxml-6.1.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:556e94a63c9b04716f8e4de2abb65775061f846e89331b6c5be79183a24f98ea", size = 5002687, upload-time = "2026-05-18T19:17:01.738Z" }, + { url = "https://files.pythonhosted.org/packages/96/1d/ab7a5c4b5a394d98a94e2d0fc67bab8297597426770dd4978370fbdaf531/lxml-6.1.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5c6bf403fbb3b3e348a561a5f4f0b9961835657981c802a1df03653eef8a9074", size = 5155099, upload-time = "2026-05-18T19:17:05.159Z" }, + { url = "https://files.pythonhosted.org/packages/d0/b1/07603bfeeb891a2596d5c2a68f7d2f70f7d11c841ebe391412c69c2857b0/lxml-6.1.1-cp310-cp310-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1dde6131244bba38a17c745836ba190bc753fd73c9291666287fd0a3fa3dcf30", size = 5057225, upload-time = "2026-05-18T19:17:08.117Z" }, + { url = "https://files.pythonhosted.org/packages/7a/16/cb391ee4b90186fa16d9ebcbe3ea96c71b8da3b0686386c8dcbcc3c67d44/lxml-6.1.1-cp310-cp310-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:98fc784c2c1440667aeedf8465bdfe10208acf0ead656a2c68627299f546b315", size = 5287643, upload-time = "2026-05-18T19:17:11.507Z" }, + { url = "https://files.pythonhosted.org/packages/eb/d6/b619717f918fd76747448fdbaee0e769edbc70e659b5b5d0112b7020b7a3/lxml-6.1.1-cp310-cp310-manylinux_2_28_i686.whl", hash = "sha256:add8cf6ddf9a65116119a28ece0f7886e30af27ba724a7594305f1d1b58a92a1", size = 5412445, upload-time = "2026-05-18T19:17:22.182Z" }, + { url = "https://files.pythonhosted.org/packages/c6/80/12bc5390ac0a3edeb579d9535e5049a5dda663438728e179d52fb319c33a/lxml-6.1.1-cp310-cp310-manylinux_2_31_armv7l.whl", hash = "sha256:cf9d57306d848218f3601fee7601fab1a327c942d56e2e97610583cb4dd74206", size = 4770864, upload-time = "2026-05-18T19:17:26.851Z" }, + { url = "https://files.pythonhosted.org/packages/0b/59/6500c09da3137f54f020e908d81cfc5ee3e8888e908fd380207afad7c2e6/lxml-6.1.1-cp310-cp310-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:88136950da4d13c318bde414ce10219931937851327f44328f2df4d2c4614067", size = 5359594, upload-time = "2026-05-18T19:17:32.527Z" }, + { url = "https://files.pythonhosted.org/packages/f2/9b/f64b4cc6b7ebcf75d95af3cde934d254b5f2f10d4163928d838d86b6eb48/lxml-6.1.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:cecdd5dfdc87b1fd87dbf81d4b037a544f47f4c744200a67013771682d67686a", size = 5107713, upload-time = "2026-05-18T19:17:04.402Z" }, + { url = "https://files.pythonhosted.org/packages/16/19/c7388ad5d3a72315d2832dc1458cbf4f2af7f2b990b606ff4876efd04511/lxml-6.1.1-cp310-cp310-musllinux_1_2_armv7l.whl", hash = "sha256:cd312b9692e831d2ffcad61eab31d91d4b4655a962e61de8fb410472cbcd37aa", size = 4803973, upload-time = "2026-05-18T19:17:06.545Z" }, + { url = "https://files.pythonhosted.org/packages/3f/22/76197f0bbf165f0b9e75be59be4997e5259cde973f12f098c1b54c7f5d60/lxml-6.1.1-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:5b7328b46d49fc9477d91ae8f6d55340347d827b7734ba3ea33faae0efef1383", size = 5349925, upload-time = "2026-05-18T19:17:09.743Z" }, + { url = "https://files.pythonhosted.org/packages/24/52/d2a0cfeccb9bcdc47c7ee05cdae5d69b48c9acf20997790a6338bb0d0b3b/lxml-6.1.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:37a58976370f36d9329d118ad0b953c5aeb9119ac9c6a4e258942a225d0573a1", size = 5309825, upload-time = "2026-05-18T19:17:13.831Z" }, + { url = "https://files.pythonhosted.org/packages/19/4a/b30944266776c2f49749ef2445aa7e78898194134b80ad776386f61b56ae/lxml-6.1.1-cp310-cp310-win32.whl", hash = "sha256:cea3f4c1af79af13cdb2da0c028111d8f8522d4f22a000c82385535f24e5cf3a", size = 3598402, upload-time = "2026-05-18T19:17:08.21Z" }, + { url = "https://files.pythonhosted.org/packages/9e/97/33691c66a4d7ec1a5a98e7c909a5b83ee45c7f7ba4cf92b1c4cf26e98079/lxml-6.1.1-cp310-cp310-win_amd64.whl", hash = "sha256:3abf332af33a74288675d936fe861fd4344da0dd6622193fbc4f2bfbb35536b5", size = 4021295, upload-time = "2026-05-18T19:17:28.638Z" }, + { url = "https://files.pythonhosted.org/packages/d0/5f/26a4dd0e12b9456ff7b12a21af5b491eb6629680d1edd73f4140fd386bcf/lxml-6.1.1-cp310-cp310-win_arm64.whl", hash = "sha256:8dadbe5b217ff35b6a8d16610dd710219b59b76d13f0e3f0d9f36786206e4485", size = 3667717, upload-time = "2026-05-19T19:22:44.474Z" }, +] + +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/4b/3541d44f3937ba468b75da9eebcae497dcf67adb65caa16760b0a6807ebb/markupsafe-3.0.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:2f981d352f04553a7171b8e44369f2af4055f888dfb147d55e42d29e29e74559", size = 11631, upload-time = "2025-09-27T18:36:05.558Z" }, + { url = "https://files.pythonhosted.org/packages/98/1b/fbd8eed11021cabd9226c37342fa6ca4e8a98d8188a8d9b66740494960e4/markupsafe-3.0.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e1c1493fb6e50ab01d20a22826e57520f1284df32f2d8601fdd90b6304601419", size = 12057, upload-time = "2025-09-27T18:36:07.165Z" }, + { url = "https://files.pythonhosted.org/packages/40/01/e560d658dc0bb8ab762670ece35281dec7b6c1b33f5fbc09ebb57a185519/markupsafe-3.0.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1ba88449deb3de88bd40044603fafffb7bc2b055d626a330323a9ed736661695", size = 22050, upload-time = "2025-09-27T18:36:08.005Z" }, + { url = "https://files.pythonhosted.org/packages/af/cd/ce6e848bbf2c32314c9b237839119c5a564a59725b53157c856e90937b7a/markupsafe-3.0.3-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f42d0984e947b8adf7dd6dde396e720934d12c506ce84eea8476409563607591", size = 20681, upload-time = "2025-09-27T18:36:08.881Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2a/b5c12c809f1c3045c4d580b035a743d12fcde53cf685dbc44660826308da/markupsafe-3.0.3-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:c0c0b3ade1c0b13b936d7970b1d37a57acde9199dc2aecc4c336773e1d86049c", size = 20705, upload-time = "2025-09-27T18:36:10.131Z" }, + { url = "https://files.pythonhosted.org/packages/cf/e3/9427a68c82728d0a88c50f890d0fc072a1484de2f3ac1ad0bfc1a7214fd5/markupsafe-3.0.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:0303439a41979d9e74d18ff5e2dd8c43ed6c6001fd40e5bf2e43f7bd9bbc523f", size = 21524, upload-time = "2025-09-27T18:36:11.324Z" }, + { url = "https://files.pythonhosted.org/packages/bc/36/23578f29e9e582a4d0278e009b38081dbe363c5e7165113fad546918a232/markupsafe-3.0.3-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:d2ee202e79d8ed691ceebae8e0486bd9a2cd4794cec4824e1c99b6f5009502f6", size = 20282, upload-time = "2025-09-27T18:36:12.573Z" }, + { url = "https://files.pythonhosted.org/packages/56/21/dca11354e756ebd03e036bd8ad58d6d7168c80ce1fe5e75218e4945cbab7/markupsafe-3.0.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:177b5253b2834fe3678cb4a5f0059808258584c559193998be2601324fdeafb1", size = 20745, upload-time = "2025-09-27T18:36:13.504Z" }, + { url = "https://files.pythonhosted.org/packages/87/99/faba9369a7ad6e4d10b6a5fbf71fa2a188fe4a593b15f0963b73859a1bbd/markupsafe-3.0.3-cp310-cp310-win32.whl", hash = "sha256:2a15a08b17dd94c53a1da0438822d70ebcd13f8c3a95abe3a9ef9f11a94830aa", size = 14571, upload-time = "2025-09-27T18:36:14.779Z" }, + { url = "https://files.pythonhosted.org/packages/d6/25/55dc3ab959917602c96985cb1253efaa4ff42f71194bddeb61eb7278b8be/markupsafe-3.0.3-cp310-cp310-win_amd64.whl", hash = "sha256:c4ffb7ebf07cfe8931028e3e4c85f0357459a3f9f9490886198848f4fa002ec8", size = 15056, upload-time = "2025-09-27T18:36:16.125Z" }, + { url = "https://files.pythonhosted.org/packages/d0/9e/0a02226640c255d1da0b8d12e24ac2aa6734da68bff14c05dd53b94a0fc3/markupsafe-3.0.3-cp310-cp310-win_arm64.whl", hash = "sha256:e2103a929dfa2fcaf9bb4e7c091983a49c9ac3b19c9061b6d5427dd7d14d81a1", size = 13932, upload-time = "2025-09-27T18:36:17.311Z" }, +] + +[[package]] +name = "matplotlib" +version = "3.10.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "contourpy" }, + { name = "cycler" }, + { name = "fonttools" }, + { name = "kiwisolver" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pillow" }, + { name = "pyparsing" }, + { name = "python-dateutil" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/63/1b/4be5be87d43d327a0cf4de1a56e86f7f84c89312452406cf122efe2839e6/matplotlib-3.10.9.tar.gz", hash = "sha256:fd66508e8c6877d98e586654b608a0456db8d7e8a546eb1e2600efd957302358", size = 34811233, upload-time = "2026-04-24T00:14:13.539Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/6f/340b04986e67aac6f66c5145ce68bf72c64bed30f92c8913499a6e6b8f99/matplotlib-3.10.9-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:77210dce9cb8153dffc967efaae990543392563d5a376d4dd8539bebcb0ed217", size = 8296625, upload-time = "2026-04-24T00:11:43.376Z" }, + { url = "https://files.pythonhosted.org/packages/bb/2f/127081eb83162053ebb9678ceac64220b93a663e0167432566e9c7c82aab/matplotlib-3.10.9-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:1e7698ac9868428e84d2c967424803b2472ff7167d9d6590d4204ed775343c3b", size = 8188790, upload-time = "2026-04-24T00:11:46.556Z" }, + { url = "https://files.pythonhosted.org/packages/fc/b7/d8bcec2626c35f96972bff656299fef4578113ea6193c8fdad324710410c/matplotlib-3.10.9-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1aa972116abb4c9d201bf245620b433726cb6856f3bef6a78f776a00f5c92d37", size = 8769389, upload-time = "2026-04-24T00:11:48.959Z" }, + { url = "https://files.pythonhosted.org/packages/12/49/b78e214a527ea732033b7f4d37f7afb504d74ba9d134bd47938230dfb8b1/matplotlib-3.10.9-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ae2f11957b27ce53497dd4d7b235c4d4f1faf383dfb39d0c5beb833bff883294", size = 9589657, upload-time = "2026-04-24T00:11:51.915Z" }, + { url = "https://files.pythonhosted.org/packages/5f/15/5246f7b43beae19c74dfee651d58d6cc8112e06f77adb4e88cc04f2e3a23/matplotlib-3.10.9-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:b049278ddce116aaa1c1377ebf58adea909132dfce0281cf7e3a1ea9fc2e2c65", size = 9651983, upload-time = "2026-04-24T00:11:54.766Z" }, + { url = "https://files.pythonhosted.org/packages/75/77/5acecfe672ba0fa1b8c0454f69ce155d1e6fc5852fa7206bf9afaf767121/matplotlib-3.10.9-cp310-cp310-win_amd64.whl", hash = "sha256:82834c3c292d24d3a8aae77cd2d20019de69d692a34a970e4fdb8d33e2ea3dda", size = 8199701, upload-time = "2026-04-24T00:11:58.389Z" }, + { url = "https://files.pythonhosted.org/packages/2c/2b/0e92ad0ac446633f928a1563db4aa8add407e1924faf0ded5b95b35afb27/matplotlib-3.10.9-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:1872fb212a05b729e649754a72d5da61d03e0554d76e80303b6f83d1d2c0552b", size = 8293058, upload-time = "2026-04-24T00:13:56.339Z" }, + { url = "https://files.pythonhosted.org/packages/4b/23/74682fd369f5299ceda438fea2a0662e6383b85c9383fb9cdfcf04713e07/matplotlib-3.10.9-pp310-pypy310_pp73-macosx_11_0_arm64.whl", hash = "sha256:985f2238880e2e69093f588f5fe2e46771747febf0649f3cf7f7b7480875317f", size = 8186627, upload-time = "2026-04-24T00:13:58.623Z" }, + { url = "https://files.pythonhosted.org/packages/ca/e8/368aab88f3c4cd8992800f31abfe0670c3e47540ba20a97e9fdbcde594b3/matplotlib-3.10.9-pp310-pypy310_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6640f75af2c6148293caa0a2b39dd806a492dd66c8a8b04035813e33d0fd2585", size = 8764117, upload-time = "2026-04-24T00:14:01.684Z" }, +] + +[[package]] +name = "mpmath" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, +] + +[[package]] +name = "mujoco" +version = "2.3.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "glfw" }, + { name = "numpy" }, + { name = "pyopengl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7c/03/7ce6078745085febd22fc4586a63016a71125031b97883d456ce1d64e5ed/mujoco-2.3.3.zip", hash = "sha256:8bd074d3c5d9d25416cf2a5b82b337a7431a6e20edbd0da7fbc05ee5255c1aaa", size = 633278, upload-time = "2023-03-20T18:23:59.89Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/21/2c/f59255b4fbd159a374c4077721bc5baac11afcc15b8a28f8c7658ac89df5/mujoco-2.3.3-2-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:7b95a0b7ae8bb9e36d04ba475a950025791c43087845235bb92bd2dd1787589a", size = 4308671, upload-time = "2023-03-20T18:09:30.955Z" }, + { url = "https://files.pythonhosted.org/packages/31/20/afc0ef5d5b9d96f3853458329f39518e638f41c0439c94ae2b97a9ab9f3a/mujoco-2.3.3-2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8b9f8e9e6d47fe60f96dc54be780a66a38d5ec2ff94d091ad54c6f87468b4b6a", size = 4177357, upload-time = "2023-03-20T18:09:34.016Z" }, + { url = "https://files.pythonhosted.org/packages/6d/da/27b0ef31aa23f64c21e7129ddf378c2548be45d87433f89e2bea4cafc811/mujoco-2.3.3-2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a9cbd3b60ac30f0b6661de050ca3e1de906b9c28ed9d084bdd708c141953247d", size = 3995762, upload-time = "2023-03-20T18:09:36.335Z" }, + { url = "https://files.pythonhosted.org/packages/6d/27/90cc9b4f88c5b797417e1fbeacb7590cd85f7e464a8ab79f60c885708e39/mujoco-2.3.3-2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9c7fd195bca86102788d86dbc2773c59c1f0ee3933b35eb0f6a86f0f2aeb5065", size = 4270849, upload-time = "2023-03-20T18:09:38.593Z" }, + { url = "https://files.pythonhosted.org/packages/cd/b3/e9119ebbbe9ea830e6c8ab7eafc0de7c82b38d6b71d2b2e38ba20e43a1b7/mujoco-2.3.3-2-cp310-cp310-win_amd64.whl", hash = "sha256:f3595e992770eff3f842cb80f7eb2b7b1b3e78995b6ecc247f98036da17ef74f", size = 3190540, upload-time = "2023-03-20T18:09:41.101Z" }, +] + +[[package]] +name = "networkx" +version = "3.4.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/1d/06475e1cd5264c0b870ea2cc6fdb3e37177c1e565c43f56ff17a10e3937f/networkx-3.4.2.tar.gz", hash = "sha256:307c3669428c5362aab27c8a1260aa8f47c4e91d3891f48be0141738d8d053e1", size = 2151368, upload-time = "2024-10-21T12:39:38.695Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b9/54/dd730b32ea14ea797530a4479b2ed46a6fb250f682a9cfb997e968bf0261/networkx-3.4.2-py3-none-any.whl", hash = "sha256:df5d4365b724cf81b8c6a7312509d0c22386097011ad1abe274afd5e9d3bbc5f", size = 1723263, upload-time = "2024-10-21T12:39:36.247Z" }, +] + +[[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/a7/94/ace0fdea5241a27d13543ee117cbc65868e82213fb31a8eb7fe9ff23f313/numpy-1.26.4-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:9ff0f4f29c51e2803569d7a51c2304de5554655a60c5d776e35b4a41413830d0", size = 20631468, upload-time = "2024-02-05T23:48:01.194Z" }, + { url = "https://files.pythonhosted.org/packages/20/f7/b24208eba89f9d1b58c1668bc6c8c4fd472b20c45573cb767f59d49fb0f6/numpy-1.26.4-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2e4ee3380d6de9c9ec04745830fd9e2eccb3e6cf790d39d7b98ffd19b0dd754a", size = 13966411, upload-time = "2024-02-05T23:48:29.038Z" }, + { url = "https://files.pythonhosted.org/packages/fc/a5/4beee6488160798683eed5bdb7eead455892c3b4e1f78d79d8d3f3b084ac/numpy-1.26.4-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d209d8969599b27ad20994c8e41936ee0964e6da07478d6c35016bc386b66ad4", size = 14219016, upload-time = "2024-02-05T23:48:54.098Z" }, + { url = "https://files.pythonhosted.org/packages/4b/d7/ecf66c1cd12dc28b4040b15ab4d17b773b87fa9d29ca16125de01adb36cd/numpy-1.26.4-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ffa75af20b44f8dba823498024771d5ac50620e6915abac414251bd971b4529f", size = 18240889, upload-time = "2024-02-05T23:49:25.361Z" }, + { url = "https://files.pythonhosted.org/packages/24/03/6f229fe3187546435c4f6f89f6d26c129d4f5bed40552899fcf1f0bf9e50/numpy-1.26.4-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:62b8e4b1e28009ef2846b4c7852046736bab361f7aeadeb6a5b89ebec3c7055a", size = 13876746, upload-time = "2024-02-05T23:49:51.983Z" }, + { url = "https://files.pythonhosted.org/packages/39/fe/39ada9b094f01f5a35486577c848fe274e374bbf8d8f472e1423a0bbd26d/numpy-1.26.4-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:a4abb4f9001ad2858e7ac189089c42178fcce737e4169dc61321660f1a96c7d2", size = 18078620, upload-time = "2024-02-05T23:50:22.515Z" }, + { url = "https://files.pythonhosted.org/packages/d5/ef/6ad11d51197aad206a9ad2286dc1aac6a378059e06e8cf22cd08ed4f20dc/numpy-1.26.4-cp310-cp310-win32.whl", hash = "sha256:bfe25acf8b437eb2a8b2d49d443800a5f18508cd811fea3181723922a8a82b07", size = 5972659, upload-time = "2024-02-05T23:50:35.834Z" }, + { url = "https://files.pythonhosted.org/packages/19/77/538f202862b9183f54108557bfda67e17603fc560c384559e769321c9d92/numpy-1.26.4-cp310-cp310-win_amd64.whl", hash = "sha256:b97fe8060236edf3662adfc2c633f56a08ae30560c56310562cb4f95500022d5", size = 15808905, upload-time = "2024-02-05T23:51:03.701Z" }, +] + +[[package]] +name = "nvidia-cublas-cu12" +version = "12.4.5.8" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ae/71/1c91302526c45ab494c23f61c7a84aa568b8c1f9d196efa5993957faf906/nvidia_cublas_cu12-12.4.5.8-py3-none-manylinux2014_x86_64.whl", hash = "sha256:2fc8da60df463fdefa81e323eef2e36489e1c94335b5358bcb38360adf75ac9b", size = 363438805, upload-time = "2024-04-03T20:57:06.025Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.4.127" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/67/42/f4f60238e8194a3106d06a058d494b18e006c10bb2b915655bd9f6ea4cb1/nvidia_cuda_cupti_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:9dec60f5ac126f7bb551c055072b69d85392b13311fcc1bcda2202d172df30fb", size = 13813957, upload-time = "2024-04-03T20:55:01.564Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.4.127" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2c/14/91ae57cd4db3f9ef7aa99f4019cfa8d54cb4caa7e00975df6467e9725a9f/nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a178759ebb095827bd30ef56598ec182b85547f1508941a3d560eb7ea1fbf338", size = 24640306, upload-time = "2024-04-03T20:56:01.463Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.4.127" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ea/27/1795d86fe88ef397885f2e580ac37628ed058a92ed2c39dc8eac3adf0619/nvidia_cuda_runtime_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:64403288fa2136ee8e467cdc9c9427e0434110899d07c779f25b5c068934faa5", size = 883737, upload-time = "2024-04-03T20:54:51.355Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.1.0.70" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/9f/fd/713452cd72343f682b1c7b9321e23829f00b842ceaedcda96e742ea0b0b3/nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl", hash = "sha256:165764f44ef8c61fcdfdfdbe769d687e06374059fbb388b6c89ecb0e28793a6f", size = 664752741, upload-time = "2024-04-22T15:24:15.253Z" }, +] + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.2.1.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/27/94/3266821f65b92b3138631e9c8e7fe1fb513804ac934485a8d05776e1dd43/nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f083fc24912aa410be21fa16d157fed2055dab1cc4b6934a0e03cba69eb242b9", size = 211459117, upload-time = "2024-04-03T20:57:40.402Z" }, +] + +[[package]] +name = "nvidia-curand-cu12" +version = "10.3.5.147" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8a/6d/44ad094874c6f1b9c654f8ed939590bdc408349f137f9b98a3a23ccec411/nvidia_curand_cu12-10.3.5.147-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a88f583d4e0bb643c49743469964103aa59f7f708d862c3ddb0fc07f851e3b8b", size = 56305206, upload-time = "2024-04-03T20:58:08.722Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.6.1.9" +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/3a/e1/5b9089a4b2a4790dfdea8b3a006052cfecff58139d5a4e34cb1a51df8d6f/nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_x86_64.whl", hash = "sha256:19e33fa442bcfd085b3086c4ebf7e8debc07cfe01e11513cc6d332fd918ac260", size = 127936057, upload-time = "2024-04-03T20:58:28.735Z" }, +] + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.3.1.170" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/db/f7/97a9ea26ed4bbbfc2d470994b8b4f338ef663be97b8f677519ac195e113d/nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_x86_64.whl", hash = "sha256:ea4f11a2904e2a8dc4b1833cc1b5181cde564edd0d5cd33e3c168eff2d1863f1", size = 207454763, upload-time = "2024-04-03T20:58:59.995Z" }, +] + +[[package]] +name = "nvidia-nccl-cu12" +version = "2.21.5" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/df/99/12cd266d6233f47d00daf3a72739872bdc10267d0383508b0b9c84a18bb6/nvidia_nccl_cu12-2.21.5-py3-none-manylinux2014_x86_64.whl", hash = "sha256:8579076d30a8c24988834445f8d633c697d42397e92ffc3f63fa26766d25e0a0", size = 188654414, upload-time = "2024-04-03T15:32:57.427Z" }, +] + +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.4.127" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ff/ff/847841bacfbefc97a00036e0fce5a0f086b640756dc38caea5e1bb002655/nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:06b3b9b25bf3f8af351d664978ca26a16d2c5127dbd53c0497e28d1fb9611d57", size = 21066810, upload-time = "2024-04-03T20:59:46.957Z" }, +] + +[[package]] +name = "nvidia-nvtx-cu12" +version = "12.4.127" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/87/20/199b8713428322a2f22b722c62b8cc278cc53dffa9705d744484b5035ee9/nvidia_nvtx_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:781e950d9b9f60d8241ccea575b32f5105a5baf4c2351cab5256a24869f12a1a", size = 99144, upload-time = "2024-04-03T20:56:12.406Z" }, +] + +[[package]] +name = "packaging" +version = "25.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a1/d4/1fc4078c65507b51b96ca8f8c3ba19e6a61c8253c72794544580a7b6c24d/packaging-25.0.tar.gz", hash = "sha256:d443872c98d677bf60f6a1f2f8c1cb748e8fe762d2bf9d3148b5599295b0fc4f", size = 165727, upload-time = "2025-04-19T11:48:59.673Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/20/12/38679034af332785aac8774540895e234f4d07f7545804097de4b666afd8/packaging-25.0-py3-none-any.whl", hash = "sha256:29572ef2b1f17581046b3a2227d5c611fb25ec70ca1ba8554b24b0e69331a484", size = 66469, upload-time = "2025-04-19T11:48:57.875Z" }, +] + +[[package]] +name = "pillow" +version = "12.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1c/3d/bb7fca845737cf9d7dbde16ed1843984665ff2e0a518f5db43e77ec540b9/pillow-12.3.0.tar.gz", hash = "sha256:3b8182a766685eaa002637e28b4ec8d6b18819a0c71f579bf0dbaa5830297cce", size = 47025035, upload-time = "2026-07-01T11:56:38.965Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/25/c2/669d88644cddb1485bd9534e63e8cf476c8e51cb3c3a1297677023505c0e/pillow-12.3.0-cp310-cp310-macosx_10_10_x86_64.whl", hash = "sha256:6c0016e7b354317c4e9e525b937ac8596c38d2d232b419529b9cd7a1cd46e39a", size = 5392418, upload-time = "2026-07-01T11:53:27.808Z" }, + { url = "https://files.pythonhosted.org/packages/6b/ba/3762f376a2948e3036488d773a146e0ae6ecc2ca03ac20e2615bd0b2ba02/pillow-12.3.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:bcc33feacfaefce60c12fd500a277533bdc02b10a19f7f6d348763d8140bbba7", size = 4785287, upload-time = "2026-07-01T11:53:29.761Z" }, + { url = "https://files.pythonhosted.org/packages/07/50/b5d688cc9c52d4482f3d5bcab6ce20bc2a74a85d2343841c907444a3be2c/pillow-12.3.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5594fc43d548a7ed94949d139aa1341b270f1863f11cfd37f5a6c8b778a6b67f", size = 6253754, upload-time = "2026-07-01T11:53:32.298Z" }, + { url = "https://files.pythonhosted.org/packages/4e/89/36f4cd76cf4baf05c50ababb976249153f18c959171c7f6ba09a6f217260/pillow-12.3.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f0606c8bf2cdefea14a43530f7657cbbb7ecf1c4222512492ef4a4434a9501ec", size = 6925605, upload-time = "2026-07-01T11:53:34.487Z" }, + { url = "https://files.pythonhosted.org/packages/eb/c0/4de58cf6633b9e3a6061ef4be6fb91fc3c90b812ece886f531e3c523d777/pillow-12.3.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:85f998ea1848bc6757289e739cfbdda3a04adfd58b02fc018ce54d754a5ce468", size = 6327788, upload-time = "2026-07-01T11:53:36.433Z" }, + { url = "https://files.pythonhosted.org/packages/87/3c/14d53682a19550dbbaf3b598f807d5457646c510805a44c7d7891cd1cd1a/pillow-12.3.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:25b9b82bb22e6e2b3cd07b39c68b7b862001226cb3dff7130d1cb914121b39ed", size = 7036288, upload-time = "2026-07-01T11:53:38.712Z" }, + { url = "https://files.pythonhosted.org/packages/38/1d/36279e3c77efe034e4cc2b0393ee74ffdb5a62391dacbf9b916154f5f0b8/pillow-12.3.0-cp310-cp310-win32.whl", hash = "sha256:37dc8f7bbb66efe481bb60defacef820c950c24713fb44962ed6aa2a50966de1", size = 6472396, upload-time = "2026-07-01T11:53:40.781Z" }, + { url = "https://files.pythonhosted.org/packages/48/7c/8fa0039574c476d7c6fa57dd7c32a130436877c6ec1e5ce1cc8ec44878c1/pillow-12.3.0-cp310-cp310-win_amd64.whl", hash = "sha256:300557495eb45ebb8aec96c2da9c4be642fbf7cd937278b4013ba894ea8eb0eb", size = 7226887, upload-time = "2026-07-01T11:53:42.764Z" }, + { url = "https://files.pythonhosted.org/packages/fa/17/e324be141d173c1c919428066c3259f21c1b8982e564e01a4a81e96dbdcf/pillow-12.3.0-cp310-cp310-win_arm64.whl", hash = "sha256:514435a37670e3e5e08f3945b68718b6ed329bb84367777e16f9f4dfe1e61a0f", size = 2568039, upload-time = "2026-07-01T11:53:45.372Z" }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + +[[package]] +name = "protobuf" +version = "7.35.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/da/01/9ef0afd7999eb9badb3a768b4aedd78c86d4c65cfaf1958ab276199e76b4/protobuf-7.35.1.tar.gz", hash = "sha256:ce115a26fe0c39a2c29973d914d327e516a6455464489fe3cd1e51a1b354f81a", size = 458717, upload-time = "2026-06-11T21:55:40.257Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/03/8aeeb7458d22546bf64b5250ca1daeb5ff757d900e8e4a7476c6f0db843e/protobuf-7.35.1-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:24f857477359a85c0c235261b8ba905fd51b2562f4a64ca1df5473f29850cbf6", size = 433226, upload-time = "2026-06-11T21:55:31.719Z" }, + { url = "https://files.pythonhosted.org/packages/37/4b/dfb89eb0e652a1ff073c39a59fb5e3a83cfe9b57a2c83fa6d78270101767/protobuf-7.35.1-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:11d6b0ec246892d85215b0a13ca6e0233cf5284b68f0ac02646427f4ff88a799", size = 328847, upload-time = "2026-06-11T21:55:34.035Z" }, + { url = "https://files.pythonhosted.org/packages/0f/58/dc12f2cd484951524af6e3382c785869b9b3fb5e52ee95ae23add53ee8f9/protobuf-7.35.1-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:b73f9489a4b8b1c9cb1f8ed951c736392592edb24b9d6819f36d2e10b171d5b4", size = 344030, upload-time = "2026-06-11T21:55:34.941Z" }, + { url = "https://files.pythonhosted.org/packages/e4/be/5b3cfe508bfab6761414ff944e3366eb13be4fd71efcd69450f89ba39f43/protobuf-7.35.1-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:74758715c53d7158fb76caf4f0cfdacc5329a4b1bb994f865d6cf302d413a1c4", size = 327130, upload-time = "2026-06-11T21:55:35.921Z" }, + { url = "https://files.pythonhosted.org/packages/d8/bc/6d6c7ba8709c85f8f2c390b2b118d6fb08a783676a572271851bf45a7d22/protobuf-7.35.1-cp310-abi3-win32.whl", hash = "sha256:353652e4efd0bca5b5fc2656abf8307ef351f0cf938c9eba09f0e09c20a25c30", size = 428945, upload-time = "2026-06-11T21:55:37.034Z" }, + { url = "https://files.pythonhosted.org/packages/0a/19/8d0cb6f20a1ef7b18f1c8986ad5783f22f84cce39c6ce9a6e645ea55192e/protobuf-7.35.1-cp310-abi3-win_amd64.whl", hash = "sha256:230a75ddfc2de4806e56696ce9640c1cdfdb6543b7cfce98d42a4c0a0e7bdb87", size = 439996, upload-time = "2026-06-11T21:55:38.123Z" }, + { url = "https://files.pythonhosted.org/packages/19/c7/5f7c636ec43e0c545e28d1f1db71990108306f7bdcb89f069ba97e428e7f/protobuf-7.35.1-py3-none-any.whl", hash = "sha256:4bc97768d8fe4ad6743c8a19403e314511ed9f6d13205b687e52421c023ac1b9", size = 171659, upload-time = "2026-06-11T21:55:39.155Z" }, +] + +[[package]] +name = "pygments" +version = "2.20.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/b2/bc9c9196916376152d655522fdcebac55e66de6603a76a02bca1b6414f6c/pygments-2.20.0.tar.gz", hash = "sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f", size = 4955991, upload-time = "2026-03-29T13:29:33.898Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" }, +] + +[[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 = "pyparsing" +version = "3.3.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f3/91/9c6ee907786a473bf81c5f53cf703ba0957b23ab84c264080fb5a450416f/pyparsing-3.3.2.tar.gz", hash = "sha256:c777f4d763f140633dcb6d8a3eda953bf7a214dc4eff598413c070bcdc117cbc", size = 6851574, upload-time = "2026-01-21T03:57:59.36Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/bd/c038d7cc38edc1aa5bf91ab8068b63d4308c66c4c8bb3cbba7dfbc049f9c/pyparsing-3.3.2-py3-none-any.whl", hash = "sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d", size = 122781, upload-time = "2026-01-21T03:57:55.912Z" }, +] + +[[package]] +name = "pyquaternion" +version = "0.9.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7d/0d/3d092aa20efaedacb89c3221a92c6491be5b28f618a2c36b52b53e7446c2/pyquaternion-0.9.9.tar.gz", hash = "sha256:b1f61af219cb2fe966b5fb79a192124f2e63a3f7a777ac3cadf2957b1a81bea8", size = 15530, upload-time = "2020-10-05T01:31:30.327Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/b3/d8482e8cacc8ea15a356efea13d22ce1c5914a9ee36622ba250523240bf2/pyquaternion-0.9.9-py3-none-any.whl", hash = "sha256:e65f6e3f7b1fdf1a9e23f82434334a1ae84f14223eee835190cd2e841f8172ec", size = 14361, upload-time = "2020-10-05T01:31:37.575Z" }, +] + +[[package]] +name = "pytest" +version = "8.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "exceptiongroup" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, + { name = "tomli" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a3/5c/00a0e072241553e1a7496d638deababa67c5058571567b92a7eaa258397c/pytest-8.4.2.tar.gz", hash = "sha256:86c0d0b93306b961d58d62a4db4879f27fe25513d4b969df351abdddb3c30e01", size = 1519618, upload-time = "2025-09-04T14:34:22.711Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a8/a4/20da314d277121d6534b3a980b29035dcd51e6744bd79075a6ce8fa4eb8d/pytest-8.4.2-py3-none-any.whl", hash = "sha256:872f880de3fc3a5bdc88a11b39c9710c3497a547cfa9320bc3c5e62fbf272e79", size = 365750, upload-time = "2025-09-04T14:34:20.226Z" }, +] + +[[package]] +name = "python-dateutil" +version = "2.9.0.post0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "six" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/c0/0c8b6ad9f17a802ee498c46e004a0eb49bc148f2fd230864601a86dcf6db/python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3", size = 342432, upload-time = "2024-03-01T18:36:20.211Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, +] + +[[package]] +name = "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 = "scipy" +version = "1.15.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/78/2f/4966032c5f8cc7e6a60f1b2e0ad686293b9474b65246b0c642e3ef3badd0/scipy-1.15.3-cp310-cp310-macosx_10_13_x86_64.whl", hash = "sha256:a345928c86d535060c9c2b25e71e87c39ab2f22fc96e9636bd74d1dbf9de448c", size = 38702770, upload-time = "2025-05-08T16:04:20.849Z" }, + { url = "https://files.pythonhosted.org/packages/a0/6e/0c3bf90fae0e910c274db43304ebe25a6b391327f3f10b5dcc638c090795/scipy-1.15.3-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:ad3432cb0f9ed87477a8d97f03b763fd1d57709f1bbde3c9369b1dff5503b253", size = 30094511, upload-time = "2025-05-08T16:04:27.103Z" }, + { url = "https://files.pythonhosted.org/packages/ea/b1/4deb37252311c1acff7f101f6453f0440794f51b6eacb1aad4459a134081/scipy-1.15.3-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:aef683a9ae6eb00728a542b796f52a5477b78252edede72b8327a886ab63293f", size = 22368151, upload-time = "2025-05-08T16:04:31.731Z" }, + { url = "https://files.pythonhosted.org/packages/38/7d/f457626e3cd3c29b3a49ca115a304cebb8cc6f31b04678f03b216899d3c6/scipy-1.15.3-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:1c832e1bd78dea67d5c16f786681b28dd695a8cb1fb90af2e27580d3d0967e92", size = 25121732, upload-time = "2025-05-08T16:04:36.596Z" }, + { url = "https://files.pythonhosted.org/packages/db/0a/92b1de4a7adc7a15dcf5bddc6e191f6f29ee663b30511ce20467ef9b82e4/scipy-1.15.3-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:263961f658ce2165bbd7b99fa5135195c3a12d9bef045345016b8b50c315cb82", size = 35547617, upload-time = "2025-05-08T16:04:43.546Z" }, + { url = "https://files.pythonhosted.org/packages/8e/6d/41991e503e51fc1134502694c5fa7a1671501a17ffa12716a4a9151af3df/scipy-1.15.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9e2abc762b0811e09a0d3258abee2d98e0c703eee49464ce0069590846f31d40", size = 37662964, upload-time = "2025-05-08T16:04:49.431Z" }, + { url = "https://files.pythonhosted.org/packages/25/e1/3df8f83cb15f3500478c889be8fb18700813b95e9e087328230b98d547ff/scipy-1.15.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:ed7284b21a7a0c8f1b6e5977ac05396c0d008b89e05498c8b7e8f4a1423bba0e", size = 37238749, upload-time = "2025-05-08T16:04:55.215Z" }, + { url = "https://files.pythonhosted.org/packages/93/3e/b3257cf446f2a3533ed7809757039016b74cd6f38271de91682aa844cfc5/scipy-1.15.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:5380741e53df2c566f4d234b100a484b420af85deb39ea35a1cc1be84ff53a5c", size = 40022383, upload-time = "2025-05-08T16:05:01.914Z" }, + { url = "https://files.pythonhosted.org/packages/d1/84/55bc4881973d3f79b479a5a2e2df61c8c9a04fcb986a213ac9c02cfb659b/scipy-1.15.3-cp310-cp310-win_amd64.whl", hash = "sha256:9d61e97b186a57350f6d6fd72640f9e99d5a4a2b8fbf4b9ee9a841eab327dc13", size = 41259201, upload-time = "2025-05-08T16:05:08.166Z" }, +] + +[[package]] +name = "setuptools" +version = "84.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6d/44/f5da03a8ef95d369145c5bb53050e7877c9f3d312e128605fd9504829143/setuptools-84.0.0.tar.gz", hash = "sha256:f4695c21257f0d9b537ec2692c941d02ee143b7cc1276941349a546573b2ef73", size = 1168449, upload-time = "2026-08-08T18:27:58.365Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/95/9c/c510029fc6ef33a6275cd2c5d3cecd6613dfd6aa401d57c54f1c18852ccf/setuptools-84.0.0-py3-none-any.whl", hash = "sha256:51a52592b3b99e102b609654876bd65f19f999935166d1352678931132b0c670", size = 818216, upload-time = "2026-08-08T18:27:56.719Z" }, +] + +[[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 = "speedtuning-sim" +version = "0.1.0" +source = { editable = "." } +dependencies = [ + { name = "dm-control" }, + { name = "mujoco" }, + { name = "numpy" }, + { name = "pyquaternion" }, +] + +[package.optional-dependencies] +evaluation = [ + { name = "matplotlib" }, +] +learned = [ + { name = "einops" }, + { name = "packaging" }, + { name = "torch" }, + { name = "torchvision" }, +] +rl = [ + { name = "torch" }, +] +test = [ + { name = "pytest" }, +] +video = [ + { name = "imageio" }, + { name = "imageio-ffmpeg" }, +] + +[package.metadata] +requires-dist = [ + { name = "dm-control", specifier = "==1.0.9" }, + { name = "einops", marker = "extra == 'learned'", specifier = ">=0.6,<0.9" }, + { name = "imageio", marker = "extra == 'video'", specifier = ">=2.31,<3" }, + { name = "imageio-ffmpeg", marker = "extra == 'video'", specifier = ">=0.4.9,<0.7" }, + { name = "matplotlib", marker = "extra == 'evaluation'", specifier = ">=3.7,<4" }, + { name = "mujoco", specifier = "==2.3.3" }, + { name = "numpy", specifier = ">=1.23,<2" }, + { name = "packaging", marker = "extra == 'learned'", specifier = ">=23,<26" }, + { name = "pyquaternion", specifier = "==0.9.9" }, + { name = "pytest", marker = "extra == 'test'", specifier = ">=8,<9" }, + { name = "torch", marker = "extra == 'learned'", specifier = "==2.5.1" }, + { name = "torch", marker = "extra == 'rl'", specifier = "==2.5.1" }, + { name = "torchvision", marker = "extra == 'learned'", specifier = "==0.20.1" }, +] +provides-extras = ["learned", "rl", "test", "video", "evaluation"] + +[[package]] +name = "sympy" +version = "1.13.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mpmath" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ca/99/5a5b6f19ff9f083671ddf7b9632028436167cd3d33e11015754e41b249a4/sympy-1.13.1.tar.gz", hash = "sha256:9cebf7e04ff162015ce31c9c6c9144daa34a93bd082f54fd8f12deca4f47515f", size = 7533040, upload-time = "2024-07-19T09:26:51.238Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b2/fe/81695a1aa331a842b582453b605175f419fe8540355886031328089d840a/sympy-1.13.1-py3-none-any.whl", hash = "sha256:db36cdc64bf61b9b24578b6f7bab1ecdd2452cf008f34faa33776680c26d66f8", size = 6189177, upload-time = "2024-07-19T09:26:48.863Z" }, +] + +[[package]] +name = "tomli" +version = "2.4.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/22/de/48c59722572767841493b26183a0d1cc411d54fd759c5607c4590b6563a6/tomli-2.4.1.tar.gz", hash = "sha256:7c7e1a961a0b2f2472c1ac5b69affa0ae1132c39adcb67aba98568702b9cc23f", size = 17543, upload-time = "2026-03-25T20:22:03.828Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7b/61/cceae43728b7de99d9b847560c262873a1f6c98202171fd5ed62640b494b/tomli-2.4.1-py3-none-any.whl", hash = "sha256:0d85819802132122da43cb86656f8d1f8c6587d54ae7dcaf30e90533028b49fe", size = 14583, upload-time = "2026-03-25T20:22:03.012Z" }, +] + +[[package]] +name = "torch" +version = "2.5.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "sympy" }, + { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/ef/834af4a885b31a0b32fff2d80e1e40f771e1566ea8ded55347502440786a/torch-2.5.1-cp310-cp310-manylinux1_x86_64.whl", hash = "sha256:71328e1bbe39d213b8721678f9dcac30dfc452a46d586f1d514a6aa0a99d4744", size = 906446312, upload-time = "2024-10-29T17:33:38.045Z" }, + { url = "https://files.pythonhosted.org/packages/69/f0/46e74e0d145f43fa506cb336eaefb2d240547e4ce1f496e442711093ab25/torch-2.5.1-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:34bfa1a852e5714cbfa17f27c49d8ce35e1b7af5608c4bc6e81392c352dbc601", size = 91919522, upload-time = "2024-10-29T17:39:08.74Z" }, + { url = "https://files.pythonhosted.org/packages/a5/13/1eb674c8efbd04d71e4a157ceba991904f633e009a584dd65dccbafbb648/torch-2.5.1-cp310-cp310-win_amd64.whl", hash = "sha256:32a037bd98a241df6c93e4c789b683335da76a2ac142c0973675b715102dc5fa", size = 203088048, upload-time = "2024-10-29T17:34:10.913Z" }, + { url = "https://files.pythonhosted.org/packages/a9/9d/e0860474ee0ff8f6ef2c50ec8f71a250f38d78a9b9df9fd241ad3397a65b/torch-2.5.1-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:23d062bf70776a3d04dbe74db950db2a5245e1ba4f27208a87f0d743b0d06e86", size = 63877046, upload-time = "2024-10-29T17:34:19.174Z" }, +] + +[[package]] +name = "torchvision" +version = "0.20.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "pillow" }, + { name = "torch" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/59/aea68d755da1451e1a0d894528a7edc9b58eb30d33e274bf21bef28dad1a/torchvision-0.20.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:4878fefb96ef293d06c27210918adc83c399d9faaf34cda5a63e129f772328f1", size = 1787552, upload-time = "2024-10-29T17:40:34.071Z" }, + { url = "https://files.pythonhosted.org/packages/a2/f6/7ff89a9f8703f623f5664afd66c8600e3f09fe188e1e0b7e6f9a8617f865/torchvision-0.20.1-cp310-cp310-manylinux1_x86_64.whl", hash = "sha256:8ffbdf8bf5b30eade22d459f5a313329eeadb20dc75efa142987b53c007098c3", size = 7238975, upload-time = "2024-10-29T17:41:03.374Z" }, + { url = "https://files.pythonhosted.org/packages/f7/ce/4c31e9b96cc4f9fec746b258d2aa35f8d1247f4f58d63f9c505ea5eb254d/torchvision-0.20.1-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:75f8a4d51a593c4bab6c9bf7d75bdd88691b00a53b07656678bc55a3a753dd73", size = 14265343, upload-time = "2024-10-29T17:40:57.799Z" }, + { url = "https://files.pythonhosted.org/packages/17/11/b5ce67715bbbec8798fb48c4a20ac28828aec1710ac01091a3eddcb8e075/torchvision-0.20.1-cp310-cp310-win_amd64.whl", hash = "sha256:22c2fa44e20eb404b85e42b22b453863a14b0927d25e550fd4f84eea97fa5b39", size = 1562413, upload-time = "2024-10-29T17:40:39.991Z" }, +] + +[[package]] +name = "tqdm" +version = "4.70.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/21/3b/6c24bec5be5e743ffd99576daa5cc077722fc7d5bbc00bd133fa0c698dc6/tqdm-4.70.0.tar.gz", hash = "sha256:55b0b0dbd97462d06ebee91e4dac24ed4d4702be82b24f07e6c1d27e08cea220", size = 795438, upload-time = "2026-07-27T11:33:15.271Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl", hash = "sha256:7f585706bfddbdebf89daac705b2dfcc16890130727d3197ca62c732b4310953", size = 80184, upload-time = "2026-07-27T11:33:13.167Z" }, +] + +[[package]] +name = "triton" +version = "3.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "filelock" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/98/29/69aa56dc0b2eb2602b553881e34243475ea2afd9699be042316842788ff5/triton-3.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6b0dd10a925263abbe9fa37dcde67a5e9b2383fc269fdf59f5657cac38c5d1d8", size = 209460013, upload-time = "2024-10-14T16:05:32.106Z" }, +] + +[[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.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, +] + +[[package]] +name = "wrapt" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2b/b0/c1f5a970721f06b85c0cd5142e0ff8fe067708abd779b0c4f4be7d61d09f/wrapt-2.3.0.tar.gz", hash = "sha256:681a2d0eefd721998f90642762b8e75c2159ec531b20ad5e437245ea7b06a107", size = 131509, upload-time = "2026-07-28T06:06:14.895Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/40/31/5822ce37ca8820c2ed35a498c67c8b37960b9cee2ba437fd32849d0a234c/wrapt-2.3.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:0bb2797048db0956348cb3058c33bc4184614f13231389cfbccc16a5d32780a7", size = 81191, upload-time = "2026-07-28T06:04:04.858Z" }, + { url = "https://files.pythonhosted.org/packages/7a/5a/3c6117938be98754578ab83f5a40d7d0ea2cd2c487dc5cd6027ee7228229/wrapt-2.3.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:ce9f398f868d2b3b27aa2ea4de79645ef9077aeeac8dfc2814b0d542c6a2b87f", size = 82255, upload-time = "2026-07-28T06:04:07.151Z" }, + { url = "https://files.pythonhosted.org/packages/a5/0f/94ae724c5087eb6054c0d63febd7094947dcf302fe058e2e0488102a872b/wrapt-2.3.0-cp310-cp310-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:ad71df7a04dd3497e9302e81f4a7c91bd401ea0e15a9df9029527900f94bee43", size = 155228, upload-time = "2026-07-28T06:04:08.272Z" }, + { url = "https://files.pythonhosted.org/packages/6c/21/1f780bba935dcf697c0c59de9be3a559bbb8e31a53ca3f25422023738432/wrapt-2.3.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fc82c2ccc8e234c844f5303d9f2984b346dcdd53e94823ce8420d2c75b4b9023", size = 157073, upload-time = "2026-07-28T06:04:09.459Z" }, + { url = "https://files.pythonhosted.org/packages/73/31/6c7799d7b6431fcd7e1b83245fb45258a2d2c3a2187fbaecb83572a72d7a/wrapt-2.3.0-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a6e19531ae33c508cea7d84a7edfda01fa86e51b8d1a93a77712c55e6e469152", size = 151594, upload-time = "2026-07-28T06:04:10.784Z" }, + { url = "https://files.pythonhosted.org/packages/ce/17/42d670dbfafd49076c6eb2b7d67633d7e1c968e39bfb11a135acb6fac67b/wrapt-2.3.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:df4ce31150bcd5d9f36f816aac3010ab4f4bf8672ac1d3b0ac7d539ec61c7c02", size = 156069, upload-time = "2026-07-28T06:04:12.316Z" }, + { url = "https://files.pythonhosted.org/packages/1f/d6/c66b4ba4eda49257c84d5c2df26118280f09ca7905aee20d0064db778d13/wrapt-2.3.0-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:e2e692bc0d63f881cf7006730a56bd4e0c2fab5dc318466942805d692b166276", size = 150930, upload-time = "2026-07-28T06:04:13.482Z" }, + { url = "https://files.pythonhosted.org/packages/c0/f2/1a3b949c0322fb27396eafd1044328c1cb0400e0b32105d75a3cd03096e7/wrapt-2.3.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:c8388ba7faf5dbf9ee106bb70d66f257629b1bd98091123e19e8a4553a319199", size = 154525, upload-time = "2026-07-28T06:04:14.698Z" }, + { url = "https://files.pythonhosted.org/packages/12/65/147563a3dfa6e830c857b93b530ebd8c0cd9d540e5914aec8f9b12880c02/wrapt-2.3.0-cp310-cp310-win32.whl", hash = "sha256:e045ff75d7d94900fc32896ed93c45ce2d2cac28c9dead582ff9a5a49d446e35", size = 77879, upload-time = "2026-07-28T06:04:16.102Z" }, + { url = "https://files.pythonhosted.org/packages/c4/eb/921405b4dc55d4f8be4c700ef120539fdd75d5fdb50d83bd257171ee18e0/wrapt-2.3.0-cp310-cp310-win_amd64.whl", hash = "sha256:b4fc96b159af0a3e0faa72475a69d66292bea72a5bed1e1aca1bffbddc3cb2b0", size = 80733, upload-time = "2026-07-28T06:04:17.43Z" }, + { url = "https://files.pythonhosted.org/packages/b6/13/75947450c5bb57795fa86384721cd52c5c4deb0879022f309501a8a85d44/wrapt-2.3.0-cp310-cp310-win_arm64.whl", hash = "sha256:1236fa25173ca964c97422470482e9011b9e3c7ed0d75798b40b3da3b0e0e760", size = 80199, upload-time = "2026-07-28T06:04:18.761Z" }, + { url = "https://files.pythonhosted.org/packages/00/39/3daf9f47be208606586de4568ba6713db53ebc8fd7a575aea1fe57983b69/wrapt-2.3.0-py3-none-any.whl", hash = "sha256:d8c7ed08477429752b8c44991f40ad7838b18332a160698740a6bfbc10d998a2", size = 61866, upload-time = "2026-07-28T06:06:12.9Z" }, +]