diff --git a/LICENSE-CC-BY-NC-4.0.md b/LICENSE-CC-BY-NC-4.0.md new file mode 100644 index 0000000..2be537d --- /dev/null +++ b/LICENSE-CC-BY-NC-4.0.md @@ -0,0 +1,159 @@ +# Attribution-NonCommercial 4.0 International + +> *Creative Commons Corporation (“Creative Commons”) is not a law firm and does not provide legal services or legal advice. Distribution of Creative Commons public licenses does not create a lawyer-client or other relationship. Creative Commons makes its licenses and related information available on an “as-is” basis. Creative Commons gives no warranties regarding its licenses, any material licensed under their terms and conditions, or any related information. Creative Commons disclaims all liability for damages resulting from their use to the fullest extent possible.* +> +> ### Using Creative Commons Public Licenses +> +> Creative Commons public licenses provide a standard set of terms and conditions that creators and other rights holders may use to share original works of authorship and other material subject to copyright and certain other rights specified in the public license below. The following considerations are for informational purposes only, are not exhaustive, and do not form part of our licenses. +> +> * __Considerations for licensors:__ Our public licenses are intended for use by those authorized to give the public permission to use material in ways otherwise restricted by copyright and certain other rights. Our licenses are irrevocable. Licensors should read and understand the terms and conditions of the license they choose before applying it. Licensors should also secure all rights necessary before applying our licenses so that the public can reuse the material as expected. Licensors should clearly mark any material not subject to the license. This includes other CC-licensed material, or material used under an exception or limitation to copyright. [More considerations for licensors](http://wiki.creativecommons.org/Considerations_for_licensors_and_licensees#Considerations_for_licensors). +> +> * __Considerations for the public:__ By using one of our public licenses, a licensor grants the public permission to use the licensed material under specified terms and conditions. If the licensor’s permission is not necessary for any reason–for example, because of any applicable exception or limitation to copyright–then that use is not regulated by the license. Our licenses grant only permissions under copyright and certain other rights that a licensor has authority to grant. Use of the licensed material may still be restricted for other reasons, including because others have copyright or other rights in the material. A licensor may make special requests, such as asking that all changes be marked or described. Although not required by our licenses, you are encouraged to respect those requests where reasonable. [More considerations for the public](http://wiki.creativecommons.org/Considerations_for_licensors_and_licensees#Considerations_for_licensees). + +## Creative Commons Attribution-NonCommercial 4.0 International Public License + +By exercising the Licensed Rights (defined below), You accept and agree to be bound by the terms and conditions of this Creative Commons Attribution-NonCommercial 4.0 International Public License ("Public License"). To the extent this Public License may be interpreted as a contract, You are granted the Licensed Rights in consideration of Your acceptance of these terms and conditions, and the Licensor grants You such rights in consideration of benefits the Licensor receives from making the Licensed Material available under these terms and conditions. + +### Section 1 – Definitions. + +a. __Adapted Material__ means material subject to Copyright and Similar Rights that is derived from or based upon the Licensed Material and in which the Licensed Material is translated, altered, arranged, transformed, or otherwise modified in a manner requiring permission under the Copyright and Similar Rights held by the Licensor. For purposes of this Public License, where the Licensed Material is a musical work, performance, or sound recording, Adapted Material is always produced where the Licensed Material is synched in timed relation with a moving image. + +b. __Adapter's License__ means the license You apply to Your Copyright and Similar Rights in Your contributions to Adapted Material in accordance with the terms and conditions of this Public License. + +c. __Copyright and Similar Rights__ means copyright and/or similar rights closely related to copyright including, without limitation, performance, broadcast, sound recording, and Sui Generis Database Rights, without regard to how the rights are labeled or categorized. For purposes of this Public License, the rights specified in Section 2(b)(1)-(2) are not Copyright and Similar Rights. + +d. __Effective Technological Measures__ means those measures that, in the absence of proper authority, may not be circumvented under laws fulfilling obligations under Article 11 of the WIPO Copyright Treaty adopted on December 20, 1996, and/or similar international agreements. + +e. __Exceptions and Limitations__ means fair use, fair dealing, and/or any other exception or limitation to Copyright and Similar Rights that applies to Your use of the Licensed Material. + +f. __Licensed Material__ means the artistic or literary work, database, or other material to which the Licensor applied this Public License. + +g. __Licensed Rights__ means the rights granted to You subject to the terms and conditions of this Public License, which are limited to all Copyright and Similar Rights that apply to Your use of the Licensed Material and that the Licensor has authority to license. + +h. __Licensor__ means the individual(s) or entity(ies) granting rights under this Public License. + +i. __NonCommercial__ means not primarily intended for or directed towards commercial advantage or monetary compensation. For purposes of this Public License, the exchange of the Licensed Material for other material subject to Copyright and Similar Rights by digital file-sharing or similar means is NonCommercial provided there is no payment of monetary compensation in connection with the exchange. + +j. __Share__ means to provide material to the public by any means or process that requires permission under the Licensed Rights, such as reproduction, public display, public performance, distribution, dissemination, communication, or importation, and to make material available to the public including in ways that members of the public may access the material from a place and at a time individually chosen by them. + +k. __Sui Generis Database Rights__ means rights other than copyright resulting from Directive 96/9/EC of the European Parliament and of the Council of 11 March 1996 on the legal protection of databases, as amended and/or succeeded, as well as other essentially equivalent rights anywhere in the world. + +l. __You__ means the individual or entity exercising the Licensed Rights under this Public License. Your has a corresponding meaning. + +### Section 2 – Scope. + +a. ___License grant.___ + + 1. Subject to the terms and conditions of this Public License, the Licensor hereby grants You a worldwide, royalty-free, non-sublicensable, non-exclusive, irrevocable license to exercise the Licensed Rights in the Licensed Material to: + + A. reproduce and Share the Licensed Material, in whole or in part, for NonCommercial purposes only; and + + B. produce, reproduce, and Share Adapted Material for NonCommercial purposes only. + + 2. __Exceptions and Limitations.__ For the avoidance of doubt, where Exceptions and Limitations apply to Your use, this Public License does not apply, and You do not need to comply with its terms and conditions. + + 3. __Term.__ The term of this Public License is specified in Section 6(a). + + 4. __Media and formats; technical modifications allowed.__ The Licensor authorizes You to exercise the Licensed Rights in all media and formats whether now known or hereafter created, and to make technical modifications necessary to do so. The Licensor waives and/or agrees not to assert any right or authority to forbid You from making technical modifications necessary to exercise the Licensed Rights, including technical modifications necessary to circumvent Effective Technological Measures. For purposes of this Public License, simply making modifications authorized by this Section 2(a)(4) never produces Adapted Material. + + 5. __Downstream recipients.__ + + A. __Offer from the Licensor – Licensed Material.__ Every recipient of the Licensed Material automatically receives an offer from the Licensor to exercise the Licensed Rights under the terms and conditions of this Public License. + + B. __No downstream restrictions.__ You may not offer or impose any additional or different terms or conditions on, or apply any Effective Technological Measures to, the Licensed Material if doing so restricts exercise of the Licensed Rights by any recipient of the Licensed Material. + + 6. __No endorsement.__ Nothing in this Public License constitutes or may be construed as permission to assert or imply that You are, or that Your use of the Licensed Material is, connected with, or sponsored, endorsed, or granted official status by, the Licensor or others designated to receive attribution as provided in Section 3(a)(1)(A)(i). + +b. ___Other rights.___ + + 1. Moral rights, such as the right of integrity, are not licensed under this Public License, nor are publicity, privacy, and/or other similar personality rights; however, to the extent possible, the Licensor waives and/or agrees not to assert any such rights held by the Licensor to the limited extent necessary to allow You to exercise the Licensed Rights, but not otherwise. + + 2. Patent and trademark rights are not licensed under this Public License. + + 3. To the extent possible, the Licensor waives any right to collect royalties from You for the exercise of the Licensed Rights, whether directly or through a collecting society under any voluntary or waivable statutory or compulsory licensing scheme. In all other cases the Licensor expressly reserves any right to collect such royalties, including when the Licensed Material is used other than for NonCommercial purposes. + +### Section 3 – License Conditions. + +Your exercise of the Licensed Rights is expressly made subject to the following conditions. + +a. ___Attribution.___ + + 1. If You Share the Licensed Material (including in modified form), You must: + + A. retain the following if it is supplied by the Licensor with the Licensed Material: + + i. identification of the creator(s) of the Licensed Material and any others designated to receive attribution, in any reasonable manner requested by the Licensor (including by pseudonym if designated); + + ii. a copyright notice; + + iii. a notice that refers to this Public License; + + iv. a notice that refers to the disclaimer of warranties; + + v. a URI or hyperlink to the Licensed Material to the extent reasonably practicable; + + B. indicate if You modified the Licensed Material and retain an indication of any previous modifications; and + + C. indicate the Licensed Material is licensed under this Public License, and include the text of, or the URI or hyperlink to, this Public License. + + 2. You may satisfy the conditions in Section 3(a)(1) in any reasonable manner based on the medium, means, and context in which You Share the Licensed Material. For example, it may be reasonable to satisfy the conditions by providing a URI or hyperlink to a resource that includes the required information. + + 3. If requested by the Licensor, You must remove any of the information required by Section 3(a)(1)(A) to the extent reasonably practicable. + + 4. If You Share Adapted Material You produce, the Adapter's License You apply must not prevent recipients of the Adapted Material from complying with this Public License. + +### Section 4 – Sui Generis Database Rights. + +Where the Licensed Rights include Sui Generis Database Rights that apply to Your use of the Licensed Material: + +a. for the avoidance of doubt, Section 2(a)(1) grants You the right to extract, reuse, reproduce, and Share all or a substantial portion of the contents of the database for NonCommercial purposes only; + +b. if You include all or a substantial portion of the database contents in a database in which You have Sui Generis Database Rights, then the database in which You have Sui Generis Database Rights (but not its individual contents) is Adapted Material; and + +c. You must comply with the conditions in Section 3(a) if You Share all or a substantial portion of the contents of the database. + +For the avoidance of doubt, this Section 4 supplements and does not replace Your obligations under this Public License where the Licensed Rights include other Copyright and Similar Rights. + +### Section 5 – Disclaimer of Warranties and Limitation of Liability. + +a. __Unless otherwise separately undertaken by the Licensor, to the extent possible, the Licensor offers the Licensed Material as-is and as-available, and makes no representations or warranties of any kind concerning the Licensed Material, whether express, implied, statutory, or other. This includes, without limitation, warranties of title, merchantability, fitness for a particular purpose, non-infringement, absence of latent or other defects, accuracy, or the presence or absence of errors, whether or not known or discoverable. Where disclaimers of warranties are not allowed in full or in part, this disclaimer may not apply to You.__ + +b. __To the extent possible, in no event will the Licensor be liable to You on any legal theory (including, without limitation, negligence) or otherwise for any direct, special, indirect, incidental, consequential, punitive, exemplary, or other losses, costs, expenses, or damages arising out of this Public License or use of the Licensed Material, even if the Licensor has been advised of the possibility of such losses, costs, expenses, or damages. Where a limitation of liability is not allowed in full or in part, this limitation may not apply to You.__ + +c. The disclaimer of warranties and limitation of liability provided above shall be interpreted in a manner that, to the extent possible, most closely approximates an absolute disclaimer and waiver of all liability. + +### Section 6 – Term and Termination. + +a. This Public License applies for the term of the Copyright and Similar Rights licensed here. However, if You fail to comply with this Public License, then Your rights under this Public License terminate automatically. + +b. Where Your right to use the Licensed Material has terminated under Section 6(a), it reinstates: + + 1. automatically as of the date the violation is cured, provided it is cured within 30 days of Your discovery of the violation; or + + 2. upon express reinstatement by the Licensor. + + For the avoidance of doubt, this Section 6(b) does not affect any right the Licensor may have to seek remedies for Your violations of this Public License. + +c. For the avoidance of doubt, the Licensor may also offer the Licensed Material under separate terms or conditions or stop distributing the Licensed Material at any time; however, doing so will not terminate this Public License. + +d. Sections 1, 5, 6, 7, and 8 survive termination of this Public License. + +### Section 7 – Other Terms and Conditions. + +a. The Licensor shall not be bound by any additional or different terms or conditions communicated by You unless expressly agreed. + +b. Any arrangements, understandings, or agreements regarding the Licensed Material not stated herein are separate from and independent of the terms and conditions of this Public License. + +### Section 8 – Interpretation. + +a. For the avoidance of doubt, this Public License does not, and shall not be interpreted to, reduce, limit, restrict, or impose conditions on any use of the Licensed Material that could lawfully be made without permission under this Public License. + +b. To the extent possible, if any provision of this Public License is deemed unenforceable, it shall be automatically reformed to the minimum extent necessary to make it enforceable. If the provision cannot be reformed, it shall be severed from this Public License without affecting the enforceability of the remaining terms and conditions. + +c. No term or condition of this Public License will be waived and no failure to comply consented to unless expressly agreed to by the Licensor. + +d. Nothing in this Public License constitutes or may be interpreted as a limitation upon, or waiver of, any privileges and immunities that apply to the Licensor or You, including from the legal processes of any jurisdiction or authority. + +> Creative Commons is not a party to its public licenses. Notwithstanding, Creative Commons may elect to apply one of its public licenses to material it publishes and in those instances will be considered the “Licensor.” Except for the limited purpose of indicating that material is shared under a Creative Commons public license or as otherwise permitted by the Creative Commons policies published at [creativecommons.org/policies](http://creativecommons.org/policies), Creative Commons does not authorize the use of the trademark “Creative Commons” or any other trademark or logo of Creative Commons without its prior written consent including, without limitation, in connection with any unauthorized modifications to any of its public licenses or any other arrangements, understandings, or agreements concerning use of licensed material. For the avoidance of doubt, this paragraph does not form part of the public licenses. +> +> Creative Commons may be contacted at creativecommons.org diff --git a/README.md b/README.md index 4c0d099..7bd2d30 100644 --- a/README.md +++ b/README.md @@ -1,18 +1,294 @@ -# 开始拉取代码 +# HUSKY: Humanoid Skateboarding System via Physics-Aware Whole-Body Control + +基于 [HUSKY](https://arxiv.org/abs/2602.03205) 思路的人形滑板全身控制实验代码：mjlab 训练、`rsl_rl` 与 MuJoCo 评测脚本。本仓库包含个人开发与 **Docker** 封装。 + +**目录：** [`src/mjlab_husky`](src/mjlab_husky) · [`rsl_rl/`](rsl_rl/) · [`dataset/`](dataset/) · [`test_scene/`](test_scene/) · [`ckpts/`](ckpts/) --- -``` -git clone http://git-test.databall.tech:3000/autocc008/robot_skating_03.git -cd robot_skating_03 +## 本地安装（Ubuntu 22.04，推荐 `uv`） -# add your files to repo -git push +```bash +curl -LsSf https://astral.sh/uv/install.sh | sh +git clone https://github.com/<你的用户名>/humanoid_skateboarding.git +cd humanoid_skateboarding +uv sync && uv pip install -e . ``` + +**（可选）LeRobot v3 导出 / 边播边录** 需要额外安装 `lerobot`（不在默认 `pyproject` 依赖里）： + +```bash +uv pip install lerobot +``` + +若安装后出现 `import torch` 报 NCCL 符号错误，可尝试： + +```bash +uv pip uninstall nvidia-nccl-cu12 +uv pip install --force-reinstall "nvidia-nccl-cu13>=2.29" +``` + --- + +## 训练 + +```bash +cd humanoid_skateboarding +uv run train Mjlab-Skater-Flat-Unitree-G1 --env.scene.num-envs 4096 +``` + +查看全部参数： + +```bash +uv run train Mjlab-Skater-Flat-Unitree-G1 --help +``` + +--- + +## 回放 `play` + +任务名固定为 **`Mjlab-Skater-Flat-Unitree-G1`**（注册在 `mjlab_husky.tasks`）。 + +### 通用 + +```bash +uv run play Mjlab-Skater-Flat-Unitree-G1 --checkpoint_file ckpts/test.pt +``` + +- **`--viewer auto`**（默认）：有 `DISPLAY` / `WAYLAND_DISPLAY` 时用 **native**，否则 **rerun**。 +- **`--viewer native`**：本机有图形界面时使用 MuJoCo 原生窗口。 +- **`--viewer rerun`**：Rerun Web Viewer（无头服务器常用）。 +- **`--viewer viser`**：Viser。 + +完整参数： + +```bash +uv run play Mjlab-Skater-Flat-Unitree-G1 --help +``` + +### 无头 OpenGL（MuJoCo 离屏相机） + +在无 `DISPLAY` 的 Linux 上，`play` 会在导入 MuJoCo 前尽量设置 **`MUJOCO_GL=egl`**（见 `mjlab_husky/mujoco_gl.py`）。若仍失败可手动指定： + +```bash +export MUJOCO_GL=egl # GPU 无头（常见） +# 或 +export MUJOCO_GL=osmesa # 纯 CPU 软件光栅（更慢） +``` + +### Rerun：端口与远程浏览器 + +Rerun 需要 **两个端口**：**Web**（默认 `8080`）+ **gRPC**（默认多为 `9876`，以终端打印为准）。 + +**浏览器与 `play` 在同一台机器**：直接打开终端里 **`http://127.0.0.1:/?url=...`** 完整链接（不要只打开无 `?url=` 的首页）。 + +**浏览器在自己电脑、`play` 在云主机**：必须在本机做 **SSH 双端口转发**（把 `user@host` 换成你的登录方式，端口与 `play` 一致）： + +```bash +ssh -N -L 8080:127.0.0.1:8080 -L 9876:127.0.0.1:9876 user@云主机IP +``` + +若使用 `~/.ssh/config` 里的 `Host` 别名（例如 `Seoul`）： + +```bash +ssh -N -L 8080:127.0.0.1:8080 -L 9876:127.0.0.1:9876 Seoul +``` + +指定密钥时： + +```bash +ssh -i ~/.ssh/你的_key -N -L 8080:127.0.0.1:8080 -L 9876:127.0.0.1:9876 ubuntu@云主机IP +``` + +**本机 8080/9876 已被占用**时，改用空闲本地端口，并同时改 `?url=` 里 gRPC 端口，例如： + +```bash +ssh -N -L 18080:127.0.0.1:18080 -L 19876:127.0.0.1:19876 user@云主机IP +``` + +云主机上 `play` 需一致： + +```bash +uv run play Mjlab-Skater-Flat-Unitree-G1 --checkpoint_file ckpts/test.pt \ + --viewer rerun \ + --rerun-web-port 18080 \ + --rerun-grpc-port 19876 +``` + +**不经 SSH、浏览器直连公网**：安全组放行 Web + gRPC 端口，并指定（示例）： + +```bash +uv run play ... --viewer rerun --rerun-connect-host <云主机公网IP> +``` + +### Rerun 常用性能参数（可选） + +```bash +uv run play Mjlab-Skater-Flat-Unitree-G1 --checkpoint_file ckpts/test.pt \ + --viewer rerun \ + --rerun-viewer-width 640 --rerun-viewer-height 360 \ + --rerun-camera-log-stride 4 --rerun-qpos-log-stride 8 \ + --rerun-camera-max-side 480 \ + --no-rerun-open-browser +``` + +说明：本项目 CLI 使用 **tyro**，布尔开关一般为 **`--xxx` / `--no-xxx`**（例如 `--lerobot-record`、`--no-rerun-open-browser`），不要写成 `--lerobot-record True`。 + +--- + +## LeRobot v3 数据（`lerobot_data/`） + +LeRobot **v3** 为 **Parquet + `meta/`**（不是 HDF5）。本仓库提供两种方式写入 **`observation.state`（qpos，float32 向量）**。 + +### 1）离线批量导出（不跑 Rerun） + +需已 `uv pip install lerobot`。 + +```bash +uv run python -m mjlab_husky.scripts.export_lerobot_qpos \ + --task-id Mjlab-Skater-Flat-Unitree-G1 \ + --checkpoint-file ckpts/test.pt \ + --out-dir lerobot_data \ + --dataset-name mjlab_husky_skater_qpos \ + --episodes 1 \ + --steps-per-episode 1000 \ + --overwrite +``` + +### 2）`play` + Rerun 同时边播边录 + +```bash +uv run play Mjlab-Skater-Flat-Unitree-G1 \ + --checkpoint_file ckpts/test.pt \ + --viewer rerun \ + --rerun-web-port 18080 \ + --rerun-grpc-port 19876 \ + --lerobot-record \ + --lerobot-out-dir lerobot_data \ + --lerobot-dataset-name mjlab_husky_live \ + --lerobot-overwrite +``` + +要点： + +- **`--lerobot-overwrite`**：每次启动会 **删除** 同名数据集目录；要 **累积** 多次运行，请 **去掉** 该参数，或换 `--lerobot-dataset-name`。 +- 默认每录满 **`--lerobot-steps-per-episode`**（默认 1000）帧会 `save_episode()` 一次；仿真里多次 `reset` **不会**自动切分，除非打开 **`--lerobot-save-on-env-reset`**。 +- 退出 `play`（如 Ctrl+C）时会 `finalize()`，避免 Parquet 不完整。 + +按仿真每次 `done -> reset` 存成一个 LeRobot episode： + +```bash +uv run play Mjlab-Skater-Flat-Unitree-G1 ... --lerobot-record --lerobot-save-on-env-reset +``` + +### 检查数据集是否可读（行数 / episode） + +```bash +uv run python -c " +from pathlib import Path +import json +info = json.loads(Path('lerobot_data/mjlab_husky_live/meta/info.json').read_text()) +print('total_episodes', info.get('total_episodes'), 'total_frames', info.get('total_frames')) +" +``` + +--- + +## Docker（推荐） + +基础环境：Ubuntu 22.04、CUDA 13、`uv` 与项目依赖。镜像 **`MUJOCO_GL=egl`**，默认 **`CMD`** 为 **Rerun** 回放（`--no-rerun-open-browser`）。 + +**构建** + +```bash +docker build -t husky-skate:latest . +``` + +**GPU 运行**（需 [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html)） + +```bash +docker run --gpus all -it --rm -p 8080:8080 -p 9876:9876 husky-skate:latest +``` + +在**宿主机浏览器**打开容器日志里打印的 **`http://127.0.0.1:8080/?url=...`**（若浏览器不在宿主机，需自行把对应端口转发到本机）。 + +**宿主机 8080/9876 已被占用**（例如已有其他容器映射）：换主机端口 + 覆盖容器内 `play` 端口，例如： + +```bash +docker run --gpus all -it --rm -p 18080:18080 -p 19876:19876 husky-skate:latest \ + uv run play Mjlab-Skater-Flat-Unitree-G1 --checkpoint_file ckpts/test.pt \ + --viewer rerun \ + --rerun-web-port 18080 \ + --rerun-grpc-port 19876 \ + --no-rerun-open-browser +``` + +**仅 CPU**（较慢） + +```bash +docker run -it --rm -p 8080:8080 -p 9876:9876 husky-skate:latest +``` + +**进入容器 Shell** + +```bash +docker run --gpus all -it --rm --entrypoint /bin/bash husky-skate:latest +``` + +**容器内训练** + +```bash +docker run --gpus all -it --rm husky-skate:latest \ + uv run train Mjlab-Skater-Flat-Unitree-G1 --env.scene.num-envs 4096 +``` + +**（可选）容器内录 LeRobot**：需先安装 `lerobot`，并把目录挂载出来，例如： + +```bash +docker run --gpus all -it --rm \ + -p 18080:18080 -p 19876:19876 \ + -v "$(pwd)/lerobot_data:/app/lerobot_data" \ + husky-skate:latest \ + bash -lc 'uv pip install lerobot && uv run play Mjlab-Skater-Flat-Unitree-G1 \ + --checkpoint_file ckpts/test.pt --viewer rerun \ + --rerun-web-port 18080 --rerun-grpc-port 19876 --no-rerun-open-browser \ + --lerobot-record --lerobot-out-dir lerobot_data --lerobot-dataset-name mjlab_docker_live \ + --lerobot-overwrite' +``` + +--- + +## PyTorch / CUDA 提示 + +若日志出现 **driver too old（如 12080）** 且 `torch.cuda.is_available()` 为 `False`，多为 **PyTorch cu13x 与当前驱动 API 不匹配**。可选： + +- 安装与驱动匹配的 **cu12x** 轮子，例如： + `uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124` +- 或升级宿主机 NVIDIA 驱动以匹配当前 PyTorch 所要求的 CUDA。 + +仿真侧 **Warp/MuJoCo** 仍可能显示 `cpu`，与 **`torch.cuda.is_available()` 为 `play` 选的 device** 一致。 + +--- + +## 轻量 MuJoCo 评测 + +```bash +bash test_scene/sim.sh your-onnx-path +``` + +| Viser | MuJoCo | +|-------|--------| +| ![](media/viser.gif) | ![](media/mjc.gif) | + +--- + +## 论文引用（原论文） + +```bibtex +@article{han2026husky, + title={HUSKY: Humanoid Skateboarding System via Physics-Aware Whole-Body Control}, + author={Jinrui Han and Dewei Wang and Chenyun Zhang and Xinzhe Liu and Ping Luo and Chenjia Bai and Xuelong Li}, + journal={arXiv preprint arXiv:2602.03205}, + year={2026} +} ``` -cd existing_repo -git remote add origin http://git-test.databall.tech:3000/autocc008/robot_skating_03.git -git branch -M main -git push -uf origin main -``` \ No newline at end of file diff --git a/autorun.sh b/autorun.sh new file mode 100755 index 0000000..e5d0d5c --- /dev/null +++ b/autorun.sh @@ -0,0 +1,9 @@ +#!/bin/bash +export PATH="$PATH:/opt/miniconda/bin" +# 加载 conda 的 shell 函数（不需要 conda init 和 source ~/.bashrc） +eval "$(conda shell.bash hook)" +conda activate pms-py312 +export PYTHONPATH="/opt/vendor/python:$CONDA_PREFIX/lib/python3.12/site-packages:/workspace/src:/workspace/rsl_rl" +python -m mjlab_husky.scripts.play Mjlab-Skater-Flat-Unitree-G1 --checkpoint_file ckpts/test.pt --viewer rerun_viser --rerun-web-port 8081 --rerun-grpc-port 9876 --viser-port 8082 + + diff --git a/ckpts/test.onnx b/ckpts/test.onnx new file mode 100644 index 0000000..bc9e13c Binary files /dev/null and b/ckpts/test.onnx differ diff --git a/ckpts/test.pt b/ckpts/test.pt new file mode 100644 index 0000000..7c337da Binary files /dev/null and b/ckpts/test.pt differ diff --git a/dataset/.DS_Store b/dataset/.DS_Store new file mode 100644 index 0000000..2935cf2 Binary files /dev/null and b/dataset/.DS_Store differ diff --git a/dataset/ref_pose/push_start_pose_b.npy b/dataset/ref_pose/push_start_pose_b.npy new file mode 100644 index 0000000..4f3de4f Binary files /dev/null and b/dataset/ref_pose/push_start_pose_b.npy differ diff --git a/dataset/ref_pose/steer_start_pose_b.npy b/dataset/ref_pose/steer_start_pose_b.npy new file mode 100644 index 0000000..40247e5 Binary files /dev/null and b/dataset/ref_pose/steer_start_pose_b.npy differ diff --git a/dataset/skate_push/human_push_1.npy b/dataset/skate_push/human_push_1.npy new file mode 100644 index 0000000..d4d8b9e Binary files /dev/null and b/dataset/skate_push/human_push_1.npy differ diff --git a/dataset/skate_push/human_push_2.npy b/dataset/skate_push/human_push_2.npy new file mode 100644 index 0000000..d849623 Binary files /dev/null and b/dataset/skate_push/human_push_2.npy differ diff --git a/lerobot_data/.DS_Store b/lerobot_data/.DS_Store new file mode 100644 index 0000000..829aa5f Binary files /dev/null and b/lerobot_data/.DS_Store differ diff --git a/lerobot_data/mjlab_husky_live/.DS_Store b/lerobot_data/mjlab_husky_live/.DS_Store new file mode 100644 index 0000000..ba38aae Binary files /dev/null and b/lerobot_data/mjlab_husky_live/.DS_Store differ diff --git a/lerobot_data/mjlab_husky_live/data/chunk-000/file-000.parquet b/lerobot_data/mjlab_husky_live/data/chunk-000/file-000.parquet new file mode 100644 index 0000000..8a65ac4 Binary files /dev/null and b/lerobot_data/mjlab_husky_live/data/chunk-000/file-000.parquet differ diff --git a/lerobot_data/mjlab_husky_live/meta/episodes/chunk-000/file-000.parquet b/lerobot_data/mjlab_husky_live/meta/episodes/chunk-000/file-000.parquet new file mode 100644 index 0000000..885f6a9 Binary files /dev/null and b/lerobot_data/mjlab_husky_live/meta/episodes/chunk-000/file-000.parquet differ diff --git a/lerobot_data/mjlab_husky_live/meta/info.json b/lerobot_data/mjlab_husky_live/meta/info.json new file mode 100644 index 0000000..2863d09 --- /dev/null +++ b/lerobot_data/mjlab_husky_live/meta/info.json @@ -0,0 +1,60 @@ +{ + "codebase_version": "v3.0", + "robot_type": "mjlab_husky", + "total_episodes": 85, + "total_frames": 85000, + "total_tasks": 1, + "chunks_size": 1000, + "data_files_size_in_mb": 100, + "video_files_size_in_mb": 200, + "fps": 50, + "splits": { + "train": "0:85" + }, + "data_path": "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet", + "video_path": null, + "features": { + "observation.state": { + "dtype": "float32", + "shape": [ + 44 + ], + "names": null + }, + "timestamp": { + "dtype": "float32", + "shape": [ + 1 + ], + "names": null + }, + "frame_index": { + "dtype": "int64", + "shape": [ + 1 + ], + "names": null + }, + "episode_index": { + "dtype": "int64", + "shape": [ + 1 + ], + "names": null + }, + "index": { + "dtype": "int64", + "shape": [ + 1 + ], + "names": null + }, + "task_index": { + "dtype": "int64", + "shape": [ + 1 + ], + "names": null + } + } +} \ No newline at end of file diff --git a/lerobot_data/mjlab_husky_live/meta/stats.json b/lerobot_data/mjlab_husky_live/meta/stats.json new file mode 100644 index 0000000..0ede314 --- /dev/null +++ b/lerobot_data/mjlab_husky_live/meta/stats.json @@ -0,0 +1,581 @@ +{ + "frame_index": { + "min": [ + 0 + ], + "max": [ + 999 + ], + "mean": [ + 499.5 + ], + "std": [ + 288.6749902572095 + ], + "count": [ + 85000 + ], + "q01": [ + 9.190799999901826 + ], + "q10": [ + 99.1007999999198 + ], + "q50": [ + 499.1003999999992 + ], + "q90": [ + 899.1000000000806 + ], + "q99": [ + 989.0100000000987 + ] + }, + "episode_index": { + "min": [ + 0 + ], + "max": [ + 84 + ], + "mean": [ + 42.0 + ], + "std": [ + 24.53568829277059 + ], + "count": [ + 85000 + ], + "q01": [ + 42.0 + ], + "q10": [ + 42.00000000000001 + ], + "q50": [ + 42.000000000000014 + ], + "q90": [ + 42.00000000000003 + ], + "q99": [ + 42.00000000000003 + ] + }, + "observation.state": { + "min": [ + -0.07113656401634216, + -70.33460998535156, + 0.0829688161611557, + 0.10346301645040512, + -0.5429396033287048, + -0.7814562320709229, + -0.9942379593849182, + -1.3800686597824097, + -0.36666902899742126, + -0.5998758673667908, + -0.12973685562610626, + -1.0291401147842407, + -0.30399635434150696, + -1.6473090648651123, + -0.854110062122345, + -0.3959100544452667, + 0.12965033948421478, + -0.9802296161651611, + -0.2205560803413391, + -0.6682753562927246, + -0.6041809916496277, + -0.5346789956092834, + -0.9392914175987244, + 0.12718690931797028, + -0.630613386631012, + -0.13077163696289062, + -1.645430564880371, + -1.463309407234192, + -0.6851853728294373, + -0.019867030903697014, + 0.0, + -70.32968139648438, + 0.0039673116989433765, + -0.2486332356929779, + -0.7575073838233948, + -0.09620650112628937, + -0.9997463822364807, + -0.2166292518377304, + -0.12168516218662262, + 0.0, + 0.0, + -0.12534703314304352, + 0.0, + 0.0 + ], + "max": [ + 75.91880798339844, + 6.302745819091797, + 0.8795564770698547, + 1.0, + 0.27965977787971497, + 0.7949097156524658, + 0.7174645066261292, + 0.4639802575111389, + 0.6382898688316345, + 0.6904199123382568, + 2.7510406970977783, + 0.44586628675460815, + 0.22848469018936157, + 0.20644104480743408, + 0.3050974905490875, + 0.8214103579521179, + 2.1551849842071533, + 0.5406686067581177, + 0.30794599652290344, + 0.9843693375587463, + 0.5602899789810181, + 0.6002522706985474, + 1.3705785274505615, + 2.2902050018310547, + 0.58452969789505, + 2.0403666496276855, + 0.9998741149902344, + -0.04809024930000305, + 0.9831812977790833, + 1.8275797367095947, + 75.97547149658203, + 6.415305137634277, + 0.17043574154376984, + 1.0, + 0.00961267203092575, + 0.7573930025100708, + 0.24439620971679688, + 0.13883520662784576, + 0.14782053232192993, + 3994.871337890625, + 3002.75146484375, + 0.15178228914737701, + 4508.38818359375, + 5436.13671875 + ], + "mean": [ + 23.25296789302546, + -14.87063912240898, + 0.7819512668777915, + 0.865570327464272, + 0.01820748194586485, + 0.06374942751341114, + -0.1310197852771072, + -0.3828951469677337, + 0.1711113832452718, + 0.03854638235369588, + 0.6976036687107648, + -0.3548905068679767, + -0.020054900660248034, + -0.779085868334069, + -0.06996407535146264, + 0.02570628469727714, + 0.9095189350492814, + -0.3067187977297342, + 0.01838621137242573, + -0.1969201212100041, + -0.051224787414813525, + 0.1772239558035782, + -0.17888976724489644, + 0.6842897481778086, + -0.23786781745798447, + 0.5855691568925977, + 0.04866229006472756, + -0.5830464264925787, + 0.07088629282490515, + 0.7353050638647644, + 23.28549128586755, + -14.991129989364966, + 0.09786677574191023, + 0.8207214686085268, + -0.014800452884902131, + 0.012796162317574593, + -0.3960311810705153, + -0.004298211561913111, + 0.004420292135738932, + 1144.7475955318005, + 1128.1047784496752, + 0.005849475387538619, + 1147.7585460326252, + 1155.0732623577117 + ], + "std": [ + 15.179133948547573, + 15.20011979901754, + 0.1454885092628437, + 0.16651873814941476, + 0.10197849704538488, + 0.1578681255507523, + 0.40764685053480554, + 0.2708406675734598, + 0.1379569269283228, + 0.13227254048456757, + 0.2890878111557964, + 0.19052926227978695, + 0.08294386306895925, + 0.3608905403828366, + 0.14723826734479117, + 0.13379134266282225, + 0.38752975069890283, + 0.22396932802488995, + 0.0868930687538448, + 0.14074262142096175, + 0.1886161113878726, + 0.11362784145745387, + 0.2621380919964973, + 0.2556253485331334, + 0.1126427040214631, + 0.3526049895335999, + 0.3032407931542005, + 0.11991412831623412, + 0.3064957559934792, + 0.3248620681656251, + 15.189723680756549, + 15.223858376740733, + 0.012404234495425687, + 0.22559996637236646, + 0.1049420524852779, + 0.09063896603938991, + 0.31475311295102665, + 0.01805742472878103, + 0.018927681317975016, + 763.1105549460016, + 734.9102254844961, + 0.025587495248672498, + 777.1573189417905, + 805.3248523580429 + ], + "count": [ + 85000 + ], + "q01": [ + 15.518811176371708, + -21.800787278659204, + 0.6495954289599476, + 0.6718457682167782, + -0.08911537257755264, + -0.09254132599076804, + -0.5442103528143726, + -0.8142312556505203, + -0.2746849088984377, + -0.37409252572026064, + 0.30180432632857673, + -0.761493222503101, + -0.18507926103823322, + -1.3567008463775416, + -0.5036703249110895, + -0.28734303055440685, + 0.28153787482213344, + -0.6747250879512114, + -0.1396642809843316, + -0.4284897106535294, + -0.42121099120915334, + -0.051688335814020214, + -0.7924406105542886, + 0.31562799811363224, + -0.37630962519084704, + 0.07698490332592935, + -0.3561562091112136, + -0.8444881160119001, + -0.3497024485512692, + 0.22996248002437986, + 15.576270143400837, + -21.96829468889173, + 0.09735096952713591, + 0.7123482675212154, + -0.018511691378479748, + 0.006816049990297945, + -0.5658663816106864, + -0.04238463985989345, + -0.03602184447898145, + 762.0721309375697, + 753.0855775318535, + -0.0461389991139774, + 763.0219338896693, + 766.1330671246137 + ], + "q10": [ + 17.13613685887827, + -20.258861770384623, + 0.7103266863599536, + 0.7094701728107882, + -0.060989131895872205, + -0.029106133983091073, + -0.504221053586579, + -0.6624676102662788, + 0.027281049858121312, + -0.030302501289641702, + 0.4506148300529305, + -0.6142739851392953, + -0.11152195437164868, + -1.1864144497536786, + -0.22565803413881974, + -0.13463749580289525, + 0.42177760002733145, + -0.537168428028581, + -0.0704819829193603, + -0.3211724682765849, + -0.31891577460017867, + 0.05491297600119163, + -0.5311642064492241, + 0.5432212542085089, + -0.3212178046207084, + 0.30091702385342545, + -0.13884038798072756, + -0.7119056475773254, + -0.21982713513540358, + 0.3988324834181499, + 17.1882684046792, + -20.450890284314816, + 0.09739910307028207, + 0.7412608709828881, + -0.01803065373541732, + 0.007440648461717085, + -0.5365755949991055, + -0.025924483710719686, + -0.013262635344442894, + 836.1614272927652, + 826.3365670006482, + -0.018764696247214368, + 837.2647427414483, + 841.0888403746054 + ], + "q50": [ + 23.762409283764843, + -14.563460296239693, + 0.7827337545593831, + 0.8979291666198453, + 0.011350150855273758, + 0.06532513037080345, + -0.12206655524496907, + -0.4096627717294724, + 0.18403609829247833, + 0.033636638799904965, + 0.6160347013881324, + -0.33975082198066753, + -0.018045310162279112, + -0.8163800627140383, + -0.09180862832401136, + 0.018329558631927374, + 0.9480252705071734, + -0.3165479713502877, + -0.004242551547214943, + -0.1876591242178994, + -0.002054170142044198, + 0.1841658829086471, + -0.1383427245692955, + 0.7028915360612112, + -0.2458919337199281, + 0.5460589266844458, + 0.0677177166342812, + -0.5704874180267574, + 0.17712665092709504, + 0.6865986548212377, + 23.765883293045736, + -14.70968608737983, + 0.0974536156380619, + 0.8232050776237005, + -0.0177864022127682, + 0.015337733944697386, + -0.3919809371150244, + -0.0017066563739786982, + 0.0019141628737798318, + 1146.1196838017734, + 1129.962561927736, + 0.0012937257823444764, + 1148.892501481066, + 1155.902512883781 + ], + "q90": [ + 28.635734952869846, + -10.041157697359191, + 0.8461450801716657, + 0.9661417386890236, + 0.10726293688968701, + 0.15083795809839803, + 0.2146376604382313, + -0.06311925098884336, + 0.3076535279334336, + 0.1550071704183138, + 1.0939077803548265, + -0.15099184687477152, + 0.06143217126236244, + -0.373247745245278, + 0.09912829460511562, + 0.17111820184132637, + 1.333676097759791, + -0.060704099186905884, + 0.12421973123064842, + -0.086594612443564, + 0.11381211995410652, + 0.2819234085433624, + 0.0842839852017181, + 0.8042130172922879, + -0.1492722455033644, + 0.907436110787069, + 0.19406254114557048, + -0.473296269342011, + 0.31749770715945097, + 1.1360847358593222, + 28.65005411403944, + -10.114658482182755, + 0.09869820489690093, + 0.8902162918360124, + -0.008530837538709927, + 0.015510477181937963, + -0.2611029862148012, + 0.012261228386334136, + 0.026974081897148575, + 1453.8357671251485, + 1429.161401413135, + 0.03886226614772836, + 1459.450008259809, + 1470.675159581278 + ], + "q99": [ + 29.426830072364876, + -8.898120419735253, + 0.8527026614373789, + 0.9813330734009837, + 0.14279720059212514, + 0.19780941931640406, + 0.31998690711516964, + 0.10990977427538703, + 0.34699264243537326, + 0.22542898462099187, + 1.2790248403157376, + -0.0044630865840350874, + 0.1601653475971783, + -0.23684657999113493, + 0.14029283857395664, + 0.2544758809367996, + 1.6298833114916786, + 0.0403672861702302, + 0.16848280955763428, + -0.007583600633284624, + 0.20617429927775735, + 0.3594991666429182, + 0.20428104321746265, + 0.9497321564037832, + 0.044834615903742187, + 1.179834731440334, + 0.3116965556845947, + -0.38883973184753867, + 0.40519437930163216, + 1.2721450699430548, + 29.490910510154002, + -8.960652529571895, + 0.09934294331161415, + 0.9097403214357518, + -0.008158426082638257, + 0.016114267153009183, + -0.21360506679982902, + 0.03291656909620061, + 0.043081196781028716, + 1524.6454015946413, + 1498.215975788589, + 0.059730523218585053, + 1530.8834143921215, + 1543.3855692444952 + ] + }, + "task_index": { + "min": [ + 0 + ], + "max": [ + 0 + ], + "mean": [ + 0.0 + ], + "std": [ + 0.0 + ], + "count": [ + 85000 + ], + "q01": [ + 3.9999999999994176e-16 + ], + "q10": [ + 3.999999999999417e-15 + ], + "q50": [ + 1.999999999999706e-14 + ], + "q90": [ + 3.599999999999482e-14 + ], + "q99": [ + 3.959999999999426e-14 + ] + }, + "index": { + "min": [ + 0 + ], + "max": [ + 84999 + ], + "mean": [ + 42499.5 + ], + "std": [ + 24537.38643886101 + ], + "count": [ + 85000 + ], + "q01": [ + 42009.19079999987 + ], + "q10": [ + 42099.100799999884 + ], + "q50": [ + 42499.100400000025 + ], + "q90": [ + 42899.10000000006 + ], + "q99": [ + 42989.01000000005 + ] + }, + "timestamp": { + "min": [ + 0.0 + ], + "max": [ + 19.98 + ], + "mean": [ + 9.99 + ], + "std": [ + 5.773499805144186 + ], + "count": [ + 85000 + ], + "q01": [ + 0.1838159999018404 + ], + "q10": [ + 1.9820159999198408 + ], + "q50": [ + 9.982007999999915 + ], + "q90": [ + 17.982000000080028 + ], + "q99": [ + 19.780200000097988 + ] + } +} \ No newline at end of file diff --git a/lerobot_data/mjlab_husky_live/meta/tasks.parquet b/lerobot_data/mjlab_husky_live/meta/tasks.parquet new file mode 100644 index 0000000..9a5f576 Binary files /dev/null and b/lerobot_data/mjlab_husky_live/meta/tasks.parquet differ diff --git a/media/method.jpg b/media/method.jpg new file mode 100644 index 0000000..263d0a7 Binary files /dev/null and b/media/method.jpg differ diff --git a/media/mjc.gif b/media/mjc.gif new file mode 100644 index 0000000..eccdf0a Binary files /dev/null and b/media/mjc.gif differ diff --git a/media/teaser.jpg b/media/teaser.jpg new file mode 100644 index 0000000..e2474fa Binary files /dev/null and b/media/teaser.jpg differ diff --git a/media/viser.gif b/media/viser.gif new file mode 100644 index 0000000..2b742f8 Binary files /dev/null and b/media/viser.gif differ diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..5fc73a1 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,42 @@ +[build-system] +requires = ["uv_build>=0.8.19,<0.9.0"] +build-backend = "uv_build" + +[project] +name = "mjlab_husky" +version = "0.1.0" +license = "Apache-2.0" +description = "Humanoid Skateboarding System via Physics-Aware Whole-Body Control" +readme = "README.md" +requires-python = ">=3.12,<3.14" +dependencies = [ + "mjlab", + "onnxruntime", + "pybullet", + "rsl-rl-lib", + "pynput", + "rerun-sdk" +] + +[project.entry-points."mjlab.tasks"] +mjlab_husky = "mjlab_husky.tasks" + +[project.scripts] +train = "mjlab_husky.scripts.train:main" +play = "mjlab_husky.scripts.play:main" +rerun-lerobot = "mjlab_husky.scripts.rerun_lerobot_dataset:main" +play-lerobot-rerun = "mjlab_husky.scripts.play_lerobot_rerun:main" + +[tool.ruff] +src = ["src"] # Helpful for recognizing first-party imports. +indent-width = 4 + +[tool.uv] +# py.mujoco.org 上的 3.7.0.dev* 预发布包会被撤下，wheel 404；强制使用 PyPI 稳定版。 +override-dependencies = ["mujoco>=3.8.0,<3.9"] +# `mjlab` 声明的 mujoco-warp git rev（1dc288c）依赖 mjENBL_MULTICCD，PyPI 的 mujoco 3.8.x 尚无该枚举。 +# `uv.lock` 将 mujoco-warp 固定为上游标签 v3.8.0（与 mujoco 3.8.0 对齐）；升级 mjlab / 跑 `uv lock` 后若冲突请复查该包。 + +[tool.uv.sources] +mjlab = { git = "https://github.com/mujocolab/mjlab.git", rev = "13212ad" } +rsl-rl-lib = { path = "rsl_rl" } diff --git a/rsl_rl/.DS_Store b/rsl_rl/.DS_Store new file mode 100644 index 0000000..767bbff Binary files /dev/null and b/rsl_rl/.DS_Store differ diff --git a/rsl_rl/__init__.py b/rsl_rl/__init__.py new file mode 100644 index 0000000..ebc2e20 --- /dev/null +++ b/rsl_rl/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Main module for the rsl_rl package.""" diff --git a/rsl_rl/__pycache__/__init__.cpython-312.pyc b/rsl_rl/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..00a46ee Binary files /dev/null and b/rsl_rl/__pycache__/__init__.cpython-312.pyc differ diff --git a/rsl_rl/algorithms/__init__.py b/rsl_rl/algorithms/__init__.py new file mode 100644 index 0000000..0ef9163 --- /dev/null +++ b/rsl_rl/algorithms/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Implementation of different RL agents.""" + +from .distillation import Distillation +from .ppo import PPO +from .amp_ppo import AMP_PPO +__all__ = ["PPO", "Distillation", "AMP_PPO"] diff --git a/rsl_rl/algorithms/__pycache__/__init__.cpython-312.pyc b/rsl_rl/algorithms/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..dfd5df6 Binary files /dev/null and b/rsl_rl/algorithms/__pycache__/__init__.cpython-312.pyc differ diff --git a/rsl_rl/algorithms/__pycache__/amp_ppo.cpython-312.pyc b/rsl_rl/algorithms/__pycache__/amp_ppo.cpython-312.pyc new file mode 100644 index 0000000..13a855e Binary files /dev/null and b/rsl_rl/algorithms/__pycache__/amp_ppo.cpython-312.pyc differ diff --git a/rsl_rl/algorithms/__pycache__/distillation.cpython-312.pyc b/rsl_rl/algorithms/__pycache__/distillation.cpython-312.pyc new file mode 100644 index 0000000..bda116d Binary files /dev/null and b/rsl_rl/algorithms/__pycache__/distillation.cpython-312.pyc differ diff --git a/rsl_rl/algorithms/__pycache__/ppo.cpython-312.pyc b/rsl_rl/algorithms/__pycache__/ppo.cpython-312.pyc new file mode 100644 index 0000000..703de42 Binary files /dev/null and b/rsl_rl/algorithms/__pycache__/ppo.cpython-312.pyc differ diff --git a/rsl_rl/algorithms/amp_ppo.py b/rsl_rl/algorithms/amp_ppo.py new file mode 100644 index 0000000..a2164eb --- /dev/null +++ b/rsl_rl/algorithms/amp_ppo.py @@ -0,0 +1,571 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations +from torch._tensor import Tensor +from torch._tensor import Tensor +from typing import Any +from copy import deepcopy + +import torch +import torch.nn as nn +import torch.optim as optim +from itertools import chain + +from rsl_rl.modules import ActorCritic +from rsl_rl.modules.rnd import RandomNetworkDistillation +from rsl_rl.storage import RolloutStorage, ReplayBufferMulti +from rsl_rl.utils import string_to_callable + + +class AMP_PPO: + """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" + + policy: ActorCritic + """The actor critic module.""" + + def __init__( + self, + policy, + discriminator, + amp_data, + amp_normalizer, + amp_num_frames=1, + amp_replay_buffer_size=100000, + num_learning_epochs=5, + num_mini_batches=4, + clip_param=0.2, + gamma=0.99, + lam=0.95, + value_loss_coef=1.0, + entropy_coef=0.01, + learning_rate=0.001, + max_grad_norm=1.0, + use_clipped_value_loss=True, + schedule="adaptive", + desired_kl=0.01, + device="cpu", + normalize_advantage_per_mini_batch=False, + # RND parameters + rnd_cfg: dict | None = None, + # Symmetry parameters + symmetry_cfg: dict | None = None, + # Distributed training parameters + multi_gpu_cfg: dict | None = None, + ): + # device-related parameters + self.device = device + self.is_multi_gpu = multi_gpu_cfg is not None + # Multi-GPU parameters + if multi_gpu_cfg is not None: + self.gpu_global_rank = multi_gpu_cfg["global_rank"] + self.gpu_world_size = multi_gpu_cfg["world_size"] + else: + self.gpu_global_rank = 0 + self.gpu_world_size = 1 + + # RND components + if rnd_cfg is not None: + # Extract parameters used in ppo + rnd_lr = rnd_cfg.pop("learning_rate", 1e-3) + # Create RND module + self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) + # Create RND optimizer + params = self.rnd.predictor.parameters() + self.rnd_optimizer = optim.Adam(params, lr=rnd_lr) + else: + self.rnd = None + self.rnd_optimizer = None + + # Symmetry components + if symmetry_cfg is not None: + # Check if symmetry is enabled + use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] + # Print that we are not using symmetry + if not use_symmetry: + print("Symmetry not used for learning. We will use it for logging instead.") + # If function is a string then resolve it to a function + if isinstance(symmetry_cfg["data_augmentation_func"], str): + symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) + # Check valid configuration + if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): + raise ValueError( + "Data augmentation enabled but the function is not callable:" + f" {symmetry_cfg['data_augmentation_func']}" + ) + # Store symmetry configuration + self.symmetry = symmetry_cfg + else: + self.symmetry = None + + ## AMP components + self.discriminator = discriminator + self.discriminator.to(self.device) + + self.amp_storage = ReplayBufferMulti(discriminator.state_dim, amp_replay_buffer_size, amp_num_frames, device) + self.amp_data = amp_data + self.amp_normalizer = amp_normalizer + + # PPO components + self.policy = policy + self.policy.to(self.device) + + # Create rollout storage + self.storage: RolloutStorage = None # type: ignore + self.transition = RolloutStorage.Transition() + self.amp_transition = RolloutStorage.Transition() + params = [ + {'params': self.policy.parameters(), 'name': 'policy'}, + ] + + params.append({ + 'params': self.discriminator.trunk.parameters(), + 'weight_decay': 10e-4, + 'name': f'amp_trunk' + }) + params.append({ + 'params': self.discriminator.amp_linear.parameters(), + 'weight_decay': 10e-2, + 'name': f'amp_head' + }) + + # Create optimizer + self.optimizer = optim.Adam(params, lr=learning_rate) + + # PPO parameters + self.clip_param = clip_param + self.num_learning_epochs = num_learning_epochs + self.num_mini_batches = num_mini_batches + self.value_loss_coef = value_loss_coef + self.entropy_coef = entropy_coef + self.gamma = gamma + self.lam = lam + self.max_grad_norm = max_grad_norm + self.use_clipped_value_loss = use_clipped_value_loss + self.desired_kl = desired_kl + self.schedule = schedule + self.learning_rate = learning_rate + self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch + + + def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): + # create rollout storage + self.storage = RolloutStorage( + training_type, + num_envs, + num_transitions_per_env, + obs, + actions_shape, + self.device, + ) + + def act(self, obs, amp_obs): + if self.policy.is_recurrent: + self.transition.hidden_states = self.policy.get_hidden_states() + # compute the actions and values + self.transition.actions = self.policy.act(obs).detach() + self.transition.values = self.policy.evaluate(obs).detach() + self.transition.actions_log_prob = self.policy.get_actions_log_prob(self.transition.actions).detach() + self.transition.action_mean = self.policy.action_mean.detach() + self.transition.action_sigma = self.policy.action_std.detach() + # need to record obs before env.step() + self.transition.observations = obs + self.amp_transition.observations = amp_obs + return self.transition.actions + + def process_env_step(self, obs, rewards, dones, extras,amp_obs, amp_obs_frames=None): + # update the normalizers + self.policy.update_normalization(obs) + if self.rnd: + self.rnd.update_normalization(obs) + + # Record the rewards and dones + # Note: we clone here because later on we bootstrap the rewards based on timeouts + self.transition.rewards = rewards.clone() + self.transition.dones = dones + + # Compute the intrinsic rewards and add to extrinsic rewards + if self.rnd: + # Compute the intrinsic rewards + self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs) + # Add intrinsic rewards to extrinsic rewards + self.transition.rewards += self.intrinsic_rewards + + # Bootstrapping on time outs + if "time_outs" in extras: + self.transition.rewards += self.gamma * torch.squeeze( + self.transition.values * extras["time_outs"].unsqueeze(1).to(self.device), 1 + ) + + if amp_obs_frames is not None: + self.amp_storage.insert(amp_obs_frames) + else: + self.amp_storage.insert(self.amp_transition.observations, amp_obs) + + # record the transition + self.storage.add_transitions(self.transition) + self.transition.clear() + self.amp_transition.clear() + self.policy.reset(dones) + + def compute_returns(self, obs): + # compute value for the last step + last_values = self.policy.evaluate(obs).detach() + self.storage.compute_returns( + last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch + ) + + def update(self): # noqa: C901 + mean_value_loss = 0 + mean_surrogate_loss = 0 + mean_entropy = 0 + mean_amp_loss = 0 + mean_grad_pen_loss = 0 + mean_policy_pred = 0 + mean_expert_pred = 0 + # -- RND loss + if self.rnd: + mean_rnd_loss = 0 + else: + mean_rnd_loss = None + # -- Symmetry loss + if self.symmetry: + mean_symmetry_loss = 0 + else: + mean_symmetry_loss = None + + # generator for mini batches + if self.policy.is_recurrent: + generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + else: + generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + + + amp_policy_generator = self.amp_storage.feed_forward_generator( + self.num_learning_epochs * self.num_mini_batches, + self.storage.num_envs * self.storage.num_transitions_per_env // self.num_mini_batches, + ) + + amp_expert_generator = self.amp_data.feed_forward_generator_23dof_multi( + self.num_learning_epochs * self.num_mini_batches, + self.storage.num_envs * self.storage.num_transitions_per_env // self.num_mini_batches, + ) + + # iterate over batches + for sample, sample_amp_policy, sample_amp_expert in zip(generator, amp_policy_generator, amp_expert_generator): + ( + obs_batch, + actions_batch, + target_values_batch, + advantages_batch, + returns_batch, + old_actions_log_prob_batch, + old_mu_batch, + old_sigma_batch, + hid_states_batch, + masks_batch, + ) = sample + + # number of augmentations per sample + # we start with 1 and increase it if we use symmetry augmentation + num_aug = 1 + # original batch size + # we assume policy group is always there and needs augmentation + original_batch_size = obs_batch.batch_size[0] + + # check if we should normalize advantages per mini batch + if self.normalize_advantage_per_mini_batch: + with torch.no_grad(): + advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8) + + # Perform symmetric augmentation + if self.symmetry and self.symmetry["use_data_augmentation"]: + # augmentation using symmetry + data_augmentation_func = self.symmetry["data_augmentation_func"] + # returned shape: [batch_size * num_aug, ...] + obs_batch, actions_batch = data_augmentation_func( + obs=obs_batch, + actions=actions_batch, + env=self.symmetry["_env"], + ) + # compute number of augmentations per sample + # we assume policy group is always there and needs augmentation + num_aug = int(obs_batch.batch_size[0] / original_batch_size) + # repeat the rest of the batch + # -- actor + old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) + # -- critic + target_values_batch = target_values_batch.repeat(num_aug, 1) + advantages_batch = advantages_batch.repeat(num_aug, 1) + returns_batch = returns_batch.repeat(num_aug, 1) + + # Recompute actions log prob and entropy for current batch of transitions + # Note: we need to do this because we updated the policy with the new parameters + # -- actor + self.policy.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) + actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch) + # -- critic + value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) + # -- entropy + # we only keep the entropy of the first augmentation (the original one) + mu_batch = self.policy.action_mean[:original_batch_size] + sigma_batch = self.policy.action_std[:original_batch_size] + entropy_batch = self.policy.entropy[:original_batch_size] + + # KL + if self.desired_kl is not None and self.schedule == "adaptive": + with torch.inference_mode(): + kl = torch.sum( + torch.log(sigma_batch / old_sigma_batch + 1.0e-5) + + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) + / (2.0 * torch.square(sigma_batch)) + - 0.5, + axis=-1, + ) + kl_mean = torch.mean(kl) + + # Reduce the KL divergence across all GPUs + if self.is_multi_gpu: + torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM) + kl_mean /= self.gpu_world_size + + # Update the learning rate + # Perform this adaptation only on the main process + # TODO: Is this needed? If KL-divergence is the "same" across all GPUs, + # then the learning rate should be the same across all GPUs. + if self.gpu_global_rank == 0: + if kl_mean > self.desired_kl * 2.0: + self.learning_rate = max(1e-5, self.learning_rate / 1.5) + elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: + self.learning_rate = min(1e-2, self.learning_rate * 1.5) + + # Update the learning rate for all GPUs + if self.is_multi_gpu: + lr_tensor = torch.tensor(self.learning_rate, device=self.device) + torch.distributed.broadcast(lr_tensor, src=0) + self.learning_rate = lr_tensor.item() + + # Update the learning rate for all parameter groups + for param_group in self.optimizer.param_groups: + param_group["lr"] = self.learning_rate + + # Surrogate loss + ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) + surrogate = -torch.squeeze(advantages_batch) * ratio + surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( + ratio, 1.0 - self.clip_param, 1.0 + self.clip_param + ) + surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() + + # Value function loss + if self.use_clipped_value_loss: + value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( + -self.clip_param, self.clip_param + ) + value_losses = (value_batch - returns_batch).pow(2) + value_losses_clipped = (value_clipped - returns_batch).pow(2) + value_loss = torch.max(value_losses, value_losses_clipped).mean() + else: + value_loss = (returns_batch - value_batch).pow(2).mean() + + loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() + + # Symmetry loss + if self.symmetry: + # obtain the symmetric actions + # if we did augmentation before then we don't need to augment again + if not self.symmetry["use_data_augmentation"]: + data_augmentation_func = self.symmetry["data_augmentation_func"] + obs_batch, _ = data_augmentation_func(obs=obs_batch, actions=None, env=self.symmetry["_env"]) + # compute number of augmentations per sample + num_aug = int(obs_batch.shape[0] / original_batch_size) + + # actions predicted by the actor for symmetrically-augmented observations + mean_actions_batch = self.policy.act_inference(obs_batch.detach().clone()) + + # compute the symmetrically augmented actions + # note: we are assuming the first augmentation is the original one. + # We do not use the action_batch from earlier since that action was sampled from the distribution. + # However, the symmetry loss is computed using the mean of the distribution. + action_mean_orig = mean_actions_batch[:original_batch_size] + _, actions_mean_symm_batch = data_augmentation_func( + obs=None, actions=action_mean_orig, env=self.symmetry["_env"] + ) + + # compute the loss (we skip the first augmentation as it is the original one) + mse_loss = torch.nn.MSELoss() + symmetry_loss = mse_loss( + mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] + ) + # add the loss to the total loss + if self.symmetry["use_mirror_loss"]: + loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss + else: + symmetry_loss = symmetry_loss.detach() + + # Random Network Distillation loss + # TODO: Move this processing to inside RND module. + if self.rnd: + # extract the rnd_state + # TODO: Check if we still need torch no grad. It is just an affine transformation. + with torch.no_grad(): + rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size]) + rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch) + # predict the embedding and the target + predicted_embedding = self.rnd.predictor(rnd_state_batch) + target_embedding = self.rnd.target(rnd_state_batch).detach() + # compute the loss as the mean squared error + mseloss = torch.nn.MSELoss() + rnd_loss = mseloss(predicted_embedding, target_embedding) + + expert_states = sample_amp_expert + policy_states = sample_amp_policy + + with torch.no_grad(): + expert_states = self.amp_normalizer.normalize_torch(expert_states.to(self.device), self.device) + policy_states = self.amp_normalizer.normalize_torch(policy_states, self.device) + + contact_phase_push = obs_batch['critic'][:, -4] + mask_push = contact_phase_push == 1. + + if mask_push.any(): + policy_d = self.discriminator(policy_states.flatten(1)) + expert_states = expert_states.to(self.device) + expert_d = self.discriminator(expert_states.flatten(1)) + + expert_loss = torch.nn.MSELoss()(expert_d, torch.ones(expert_d.size(), device=self.device)) + policy_loss = torch.nn.MSELoss()(policy_d, -1 * torch.ones(policy_d.size(), device=self.device)) + amp_loss = 0.5 * (expert_loss + policy_loss) + + # grad penalty + grad_pen_loss = self.discriminator.compute_grad_pen(expert_states, lambda_=5) + else: + amp_loss = torch.tensor(0.0, device=self.device) + grad_pen_loss = torch.tensor(0.0, device=self.device) + expert_loss = torch.tensor(0.0, device=self.device) + policy_loss = torch.tensor(0.0, device=self.device) + + loss += (amp_loss + grad_pen_loss) + self.amp_normalizer.update(policy_states.cpu().numpy()) + self.amp_normalizer.update(expert_states.cpu().numpy()) + + # Compute the gradients + # -- For PPO + self.optimizer.zero_grad() + loss.backward() + # -- For RND + if self.rnd: + self.rnd_optimizer.zero_grad() # type: ignore + rnd_loss.backward() + + # Collect gradients from all GPUs + if self.is_multi_gpu: + self.reduce_parameters() + + # Apply the gradients + # -- For PPO + nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) + self.optimizer.step() + # -- For RND + if self.rnd_optimizer: + self.rnd_optimizer.step() + + # Store the losses + mean_value_loss += value_loss.item() + mean_surrogate_loss += surrogate_loss.item() + mean_entropy += entropy_batch.mean().item() + mean_amp_loss += amp_loss.item() + mean_grad_pen_loss += grad_pen_loss.item() + mean_policy_pred += policy_loss.mean().item() + mean_expert_pred += expert_loss.mean().item() + # -- RND loss + if mean_rnd_loss is not None: + mean_rnd_loss += rnd_loss.item() + # -- Symmetry loss + if mean_symmetry_loss is not None: + mean_symmetry_loss += symmetry_loss.item() + + # -- For PPO + num_updates = self.num_learning_epochs * self.num_mini_batches + mean_value_loss /= num_updates + mean_surrogate_loss /= num_updates + mean_entropy /= num_updates + mean_amp_loss /= num_updates + mean_grad_pen_loss /= num_updates + mean_policy_pred /= num_updates + mean_expert_pred /= num_updates + + # -- For RND + if mean_rnd_loss is not None: + mean_rnd_loss /= num_updates + # -- For Symmetry + if mean_symmetry_loss is not None: + mean_symmetry_loss /= num_updates + # -- Clear the storage + self.storage.clear() + + # construct the loss dictionary + loss_dict = { + "value_function": mean_value_loss, + "surrogate": mean_surrogate_loss, + "entropy": mean_entropy, + "amp": mean_amp_loss, + "amp_grad_pen": mean_grad_pen_loss, + "amp_policy_pred": mean_policy_pred, + "amp_expert_pred": mean_expert_pred, + } + if self.rnd: + loss_dict["rnd"] = mean_rnd_loss + if self.symmetry: + loss_dict["symmetry"] = mean_symmetry_loss + + return loss_dict + + """ + Helper functions + """ + + def broadcast_parameters(self): + """Broadcast model parameters to all GPUs.""" + # obtain the model parameters on current GPU + model_params = [self.policy.state_dict()] + if self.rnd: + model_params.append(self.rnd.predictor.state_dict()) + # broadcast the model parameters + torch.distributed.broadcast_object_list(model_params, src=0) + # load the model parameters on all GPUs from source GPU + self.policy.load_state_dict(model_params[0]) + if self.rnd: + self.rnd.predictor.load_state_dict(model_params[1]) + + def reduce_parameters(self): + """Collect gradients from all GPUs and average them. + + This function is called after the backward pass to synchronize the gradients across all GPUs. + """ + # Create a tensor to store the gradients + grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] + if self.rnd: + grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None] + all_grads = torch.cat(grads) + + # Average the gradients across all GPUs + torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) + all_grads /= self.gpu_world_size + + # Get all parameters + all_params = self.policy.parameters() + if self.rnd: + all_params = chain(all_params, self.rnd.parameters()) + + # Update the gradients for all parameters with the reduced gradients + offset = 0 + for param in all_params: + if param.grad is not None: + numel = param.numel() + # copy data back from shared buffer + param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) + # update the offset for the next parameter + offset += numel diff --git a/rsl_rl/algorithms/distillation.py b/rsl_rl/algorithms/distillation.py new file mode 100644 index 0000000..3a86e00 --- /dev/null +++ b/rsl_rl/algorithms/distillation.py @@ -0,0 +1,185 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch +import torch.nn as nn + +from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent +from rsl_rl.storage import RolloutStorage +from rsl_rl.utils import resolve_optimizer + + +class Distillation: + """Distillation algorithm for training a student model to mimic a teacher model.""" + + policy: StudentTeacher | StudentTeacherRecurrent + """The student teacher model.""" + + def __init__( + self, + policy, + num_learning_epochs=1, + gradient_length=15, + learning_rate=1e-3, + max_grad_norm=None, + loss_type="mse", + optimizer="adam", + device="cpu", + # Distributed training parameters + multi_gpu_cfg: dict | None = None, + ): + # device-related parameters + self.device = device + self.is_multi_gpu = multi_gpu_cfg is not None + # Multi-GPU parameters + if multi_gpu_cfg is not None: + self.gpu_global_rank = multi_gpu_cfg["global_rank"] + self.gpu_world_size = multi_gpu_cfg["world_size"] + else: + self.gpu_global_rank = 0 + self.gpu_world_size = 1 + + # distillation components + self.policy = policy + self.policy.to(self.device) + self.storage = None # initialized later + + # initialize the optimizer + self.optimizer = resolve_optimizer(optimizer)(self.policy.parameters(), lr=learning_rate) + + # initialize the transition + self.transition = RolloutStorage.Transition() + self.last_hidden_states = None + + # distillation parameters + self.num_learning_epochs = num_learning_epochs + self.gradient_length = gradient_length + self.learning_rate = learning_rate + self.max_grad_norm = max_grad_norm + + # initialize the loss function + loss_fn_dict = { + "mse": nn.functional.mse_loss, + "huber": nn.functional.huber_loss, + } + if loss_type in loss_fn_dict: + self.loss_fn = loss_fn_dict[loss_type] + else: + raise ValueError(f"Unknown loss type: {loss_type}. Supported types are: {list(loss_fn_dict.keys())}") + + self.num_updates = 0 + + def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): + # create rollout storage + self.storage = RolloutStorage( + training_type, + num_envs, + num_transitions_per_env, + obs, + actions_shape, + self.device, + ) + + def act(self, obs): + # compute the actions + self.transition.actions = self.policy.act(obs).detach() + self.transition.privileged_actions = self.policy.evaluate(obs).detach() + # record the observations + self.transition.observations = obs + return self.transition.actions + + def process_env_step(self, obs, rewards, dones, extras): + # update the normalizers + self.policy.update_normalization(obs) + + # record the rewards and dones + self.transition.rewards = rewards + self.transition.dones = dones + # record the transition + self.storage.add_transitions(self.transition) + self.transition.clear() + self.policy.reset(dones) + + def update(self): + self.num_updates += 1 + mean_behavior_loss = 0 + loss = 0 + cnt = 0 + + for epoch in range(self.num_learning_epochs): + self.policy.reset(hidden_states=self.last_hidden_states) + self.policy.detach_hidden_states() + for obs, _, privileged_actions, dones in self.storage.generator(): + + # inference the student for gradient computation + actions = self.policy.act_inference(obs) + + # behavior cloning loss + behavior_loss = self.loss_fn(actions, privileged_actions) + + # total loss + loss = loss + behavior_loss + mean_behavior_loss += behavior_loss.item() + cnt += 1 + + # gradient step + if cnt % self.gradient_length == 0: + self.optimizer.zero_grad() + loss.backward() + if self.is_multi_gpu: + self.reduce_parameters() + if self.max_grad_norm: + nn.utils.clip_grad_norm_(self.policy.student.parameters(), self.max_grad_norm) + self.optimizer.step() + self.policy.detach_hidden_states() + loss = 0 + + # reset dones + self.policy.reset(dones.view(-1)) + self.policy.detach_hidden_states(dones.view(-1)) + + mean_behavior_loss /= cnt + self.storage.clear() + self.last_hidden_states = self.policy.get_hidden_states() + self.policy.detach_hidden_states() + + # construct the loss dictionary + loss_dict = {"behavior": mean_behavior_loss} + + return loss_dict + + """ + Helper functions + """ + + def broadcast_parameters(self): + """Broadcast model parameters to all GPUs.""" + # obtain the model parameters on current GPU + model_params = [self.policy.state_dict()] + # broadcast the model parameters + torch.distributed.broadcast_object_list(model_params, src=0) + # load the model parameters on all GPUs from source GPU + self.policy.load_state_dict(model_params[0]) + + def reduce_parameters(self): + """Collect gradients from all GPUs and average them. + + This function is called after the backward pass to synchronize the gradients across all GPUs. + """ + # Create a tensor to store the gradients + grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] + all_grads = torch.cat(grads) + # Average the gradients across all GPUs + torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) + all_grads /= self.gpu_world_size + # Update the gradients for all parameters with the reduced gradients + offset = 0 + for param in self.policy.parameters(): + if param.grad is not None: + numel = param.numel() + # copy data back from shared buffer + param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) + # update the offset for the next parameter + offset += numel diff --git a/rsl_rl/algorithms/ppo.py b/rsl_rl/algorithms/ppo.py new file mode 100644 index 0000000..6c21fc5 --- /dev/null +++ b/rsl_rl/algorithms/ppo.py @@ -0,0 +1,469 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.optim as optim +from itertools import chain + +from rsl_rl.modules import ActorCritic +from rsl_rl.modules.rnd import RandomNetworkDistillation +from rsl_rl.storage import RolloutStorage +from rsl_rl.utils import string_to_callable + + +class PPO: + """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" + + policy: ActorCritic + """The actor critic module.""" + + def __init__( + self, + policy, + num_learning_epochs=5, + num_mini_batches=4, + clip_param=0.2, + gamma=0.99, + lam=0.95, + value_loss_coef=1.0, + entropy_coef=0.01, + learning_rate=0.001, + max_grad_norm=1.0, + use_clipped_value_loss=True, + schedule="adaptive", + desired_kl=0.01, + device="cpu", + normalize_advantage_per_mini_batch=False, + # RND parameters + rnd_cfg: dict | None = None, + # Symmetry parameters + symmetry_cfg: dict | None = None, + # Distributed training parameters + multi_gpu_cfg: dict | None = None, + ): + # device-related parameters + self.device = device + self.is_multi_gpu = multi_gpu_cfg is not None + # Multi-GPU parameters + if multi_gpu_cfg is not None: + self.gpu_global_rank = multi_gpu_cfg["global_rank"] + self.gpu_world_size = multi_gpu_cfg["world_size"] + else: + self.gpu_global_rank = 0 + self.gpu_world_size = 1 + + # RND components + if rnd_cfg is not None: + # Extract parameters used in ppo + rnd_lr = rnd_cfg.pop("learning_rate", 1e-3) + # Create RND module + self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) + # Create RND optimizer + params = self.rnd.predictor.parameters() + self.rnd_optimizer = optim.Adam(params, lr=rnd_lr) + else: + self.rnd = None + self.rnd_optimizer = None + + # Symmetry components + if symmetry_cfg is not None: + # Check if symmetry is enabled + use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] + # Print that we are not using symmetry + if not use_symmetry: + print("Symmetry not used for learning. We will use it for logging instead.") + # If function is a string then resolve it to a function + if isinstance(symmetry_cfg["data_augmentation_func"], str): + symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) + # Check valid configuration + if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): + raise ValueError( + "Data augmentation enabled but the function is not callable:" + f" {symmetry_cfg['data_augmentation_func']}" + ) + # Store symmetry configuration + self.symmetry = symmetry_cfg + else: + self.symmetry = None + + # PPO components + self.policy = policy + self.policy.to(self.device) + # Create optimizer + self.optimizer = optim.Adam(self.policy.parameters(), lr=learning_rate) + # Create rollout storage + self.storage: RolloutStorage = None # type: ignore + self.transition = RolloutStorage.Transition() + + # PPO parameters + self.clip_param = clip_param + self.num_learning_epochs = num_learning_epochs + self.num_mini_batches = num_mini_batches + self.value_loss_coef = value_loss_coef + self.entropy_coef = entropy_coef + self.gamma = gamma + self.lam = lam + self.max_grad_norm = max_grad_norm + self.use_clipped_value_loss = use_clipped_value_loss + self.desired_kl = desired_kl + self.schedule = schedule + self.learning_rate = learning_rate + self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch + + def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): + # create rollout storage + self.storage = RolloutStorage( + training_type, + num_envs, + num_transitions_per_env, + obs, + actions_shape, + self.device, + ) + + def act(self, obs): + if self.policy.is_recurrent: + self.transition.hidden_states = self.policy.get_hidden_states() + # compute the actions and values + self.transition.actions = self.policy.act(obs).detach() + self.transition.values = self.policy.evaluate(obs).detach() + self.transition.actions_log_prob = self.policy.get_actions_log_prob(self.transition.actions).detach() + self.transition.action_mean = self.policy.action_mean.detach() + self.transition.action_sigma = self.policy.action_std.detach() + # need to record obs before env.step() + self.transition.observations = obs + return self.transition.actions + + def process_env_step(self, obs, rewards, dones, extras): + # update the normalizers + self.policy.update_normalization(obs) + if self.rnd: + self.rnd.update_normalization(obs) + + # Record the rewards and dones + # Note: we clone here because later on we bootstrap the rewards based on timeouts + self.transition.rewards = rewards.clone() + self.transition.dones = dones + + # Compute the intrinsic rewards and add to extrinsic rewards + if self.rnd: + # Compute the intrinsic rewards + self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs) + # Add intrinsic rewards to extrinsic rewards + self.transition.rewards += self.intrinsic_rewards + + # Bootstrapping on time outs + if "time_outs" in extras: + self.transition.rewards += self.gamma * torch.squeeze( + self.transition.values * extras["time_outs"].unsqueeze(1).to(self.device), 1 + ) + + # record the transition + self.storage.add_transitions(self.transition) + self.transition.clear() + self.policy.reset(dones) + + def compute_returns(self, obs): + # compute value for the last step + last_values = self.policy.evaluate(obs).detach() + self.storage.compute_returns( + last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch + ) + + def update(self): # noqa: C901 + mean_value_loss = 0 + mean_surrogate_loss = 0 + mean_entropy = 0 + # -- RND loss + if self.rnd: + mean_rnd_loss = 0 + else: + mean_rnd_loss = None + # -- Symmetry loss + if self.symmetry: + mean_symmetry_loss = 0 + else: + mean_symmetry_loss = None + + # generator for mini batches + if self.policy.is_recurrent: + generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + else: + generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + + # iterate over batches + for ( + obs_batch, + actions_batch, + target_values_batch, + advantages_batch, + returns_batch, + old_actions_log_prob_batch, + old_mu_batch, + old_sigma_batch, + hid_states_batch, + masks_batch, + ) in generator: + + # number of augmentations per sample + # we start with 1 and increase it if we use symmetry augmentation + num_aug = 1 + # original batch size + # we assume policy group is always there and needs augmentation + original_batch_size = obs_batch.batch_size[0] + + # check if we should normalize advantages per mini batch + if self.normalize_advantage_per_mini_batch: + with torch.no_grad(): + advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8) + + # Perform symmetric augmentation + if self.symmetry and self.symmetry["use_data_augmentation"]: + # augmentation using symmetry + data_augmentation_func = self.symmetry["data_augmentation_func"] + # returned shape: [batch_size * num_aug, ...] + obs_batch, actions_batch = data_augmentation_func( + obs=obs_batch, + actions=actions_batch, + env=self.symmetry["_env"], + ) + # compute number of augmentations per sample + # we assume policy group is always there and needs augmentation + num_aug = int(obs_batch.batch_size[0] / original_batch_size) + # repeat the rest of the batch + # -- actor + old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) + # -- critic + target_values_batch = target_values_batch.repeat(num_aug, 1) + advantages_batch = advantages_batch.repeat(num_aug, 1) + returns_batch = returns_batch.repeat(num_aug, 1) + + # Recompute actions log prob and entropy for current batch of transitions + # Note: we need to do this because we updated the policy with the new parameters + # -- actor + self.policy.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) + actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch) + # -- critic + value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) + # -- entropy + # we only keep the entropy of the first augmentation (the original one) + mu_batch = self.policy.action_mean[:original_batch_size] + sigma_batch = self.policy.action_std[:original_batch_size] + entropy_batch = self.policy.entropy[:original_batch_size] + + # KL + if self.desired_kl is not None and self.schedule == "adaptive": + with torch.inference_mode(): + kl = torch.sum( + torch.log(sigma_batch / old_sigma_batch + 1.0e-5) + + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) + / (2.0 * torch.square(sigma_batch)) + - 0.5, + axis=-1, + ) + kl_mean = torch.mean(kl) + + # Reduce the KL divergence across all GPUs + if self.is_multi_gpu: + torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM) + kl_mean /= self.gpu_world_size + + # Update the learning rate + # Perform this adaptation only on the main process + # TODO: Is this needed? If KL-divergence is the "same" across all GPUs, + # then the learning rate should be the same across all GPUs. + if self.gpu_global_rank == 0: + if kl_mean > self.desired_kl * 2.0: + self.learning_rate = max(1e-5, self.learning_rate / 1.5) + elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: + self.learning_rate = min(1e-2, self.learning_rate * 1.5) + + # Update the learning rate for all GPUs + if self.is_multi_gpu: + lr_tensor = torch.tensor(self.learning_rate, device=self.device) + torch.distributed.broadcast(lr_tensor, src=0) + self.learning_rate = lr_tensor.item() + + # Update the learning rate for all parameter groups + for param_group in self.optimizer.param_groups: + param_group["lr"] = self.learning_rate + + # Surrogate loss + ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) + surrogate = -torch.squeeze(advantages_batch) * ratio + surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( + ratio, 1.0 - self.clip_param, 1.0 + self.clip_param + ) + surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() + + # Value function loss + if self.use_clipped_value_loss: + value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( + -self.clip_param, self.clip_param + ) + value_losses = (value_batch - returns_batch).pow(2) + value_losses_clipped = (value_clipped - returns_batch).pow(2) + value_loss = torch.max(value_losses, value_losses_clipped).mean() + else: + value_loss = (returns_batch - value_batch).pow(2).mean() + + loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() + + # Symmetry loss + if self.symmetry: + # obtain the symmetric actions + # if we did augmentation before then we don't need to augment again + if not self.symmetry["use_data_augmentation"]: + data_augmentation_func = self.symmetry["data_augmentation_func"] + obs_batch, _ = data_augmentation_func(obs=obs_batch, actions=None, env=self.symmetry["_env"]) + # compute number of augmentations per sample + num_aug = int(obs_batch.shape[0] / original_batch_size) + + # actions predicted by the actor for symmetrically-augmented observations + mean_actions_batch = self.policy.act_inference(obs_batch.detach().clone()) + + # compute the symmetrically augmented actions + # note: we are assuming the first augmentation is the original one. + # We do not use the action_batch from earlier since that action was sampled from the distribution. + # However, the symmetry loss is computed using the mean of the distribution. + action_mean_orig = mean_actions_batch[:original_batch_size] + _, actions_mean_symm_batch = data_augmentation_func( + obs=None, actions=action_mean_orig, env=self.symmetry["_env"] + ) + + # compute the loss (we skip the first augmentation as it is the original one) + mse_loss = torch.nn.MSELoss() + symmetry_loss = mse_loss( + mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] + ) + # add the loss to the total loss + if self.symmetry["use_mirror_loss"]: + loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss + else: + symmetry_loss = symmetry_loss.detach() + + # Random Network Distillation loss + # TODO: Move this processing to inside RND module. + if self.rnd: + # extract the rnd_state + # TODO: Check if we still need torch no grad. It is just an affine transformation. + with torch.no_grad(): + rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size]) + rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch) + # predict the embedding and the target + predicted_embedding = self.rnd.predictor(rnd_state_batch) + target_embedding = self.rnd.target(rnd_state_batch).detach() + # compute the loss as the mean squared error + mseloss = torch.nn.MSELoss() + rnd_loss = mseloss(predicted_embedding, target_embedding) + + # Compute the gradients + # -- For PPO + self.optimizer.zero_grad() + loss.backward() + # -- For RND + if self.rnd: + self.rnd_optimizer.zero_grad() # type: ignore + rnd_loss.backward() + + # Collect gradients from all GPUs + if self.is_multi_gpu: + self.reduce_parameters() + + # Apply the gradients + # -- For PPO + nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) + self.optimizer.step() + # -- For RND + if self.rnd_optimizer: + self.rnd_optimizer.step() + + # Store the losses + mean_value_loss += value_loss.item() + mean_surrogate_loss += surrogate_loss.item() + mean_entropy += entropy_batch.mean().item() + # -- RND loss + if mean_rnd_loss is not None: + mean_rnd_loss += rnd_loss.item() + # -- Symmetry loss + if mean_symmetry_loss is not None: + mean_symmetry_loss += symmetry_loss.item() + + # -- For PPO + num_updates = self.num_learning_epochs * self.num_mini_batches + mean_value_loss /= num_updates + mean_surrogate_loss /= num_updates + mean_entropy /= num_updates + # -- For RND + if mean_rnd_loss is not None: + mean_rnd_loss /= num_updates + # -- For Symmetry + if mean_symmetry_loss is not None: + mean_symmetry_loss /= num_updates + # -- Clear the storage + self.storage.clear() + + # construct the loss dictionary + loss_dict = { + "value_function": mean_value_loss, + "surrogate": mean_surrogate_loss, + "entropy": mean_entropy, + } + if self.rnd: + loss_dict["rnd"] = mean_rnd_loss + if self.symmetry: + loss_dict["symmetry"] = mean_symmetry_loss + + return loss_dict + + """ + Helper functions + """ + + def broadcast_parameters(self): + """Broadcast model parameters to all GPUs.""" + # obtain the model parameters on current GPU + model_params = [self.policy.state_dict()] + if self.rnd: + model_params.append(self.rnd.predictor.state_dict()) + # broadcast the model parameters + torch.distributed.broadcast_object_list(model_params, src=0) + # load the model parameters on all GPUs from source GPU + self.policy.load_state_dict(model_params[0]) + if self.rnd: + self.rnd.predictor.load_state_dict(model_params[1]) + + def reduce_parameters(self): + """Collect gradients from all GPUs and average them. + + This function is called after the backward pass to synchronize the gradients across all GPUs. + """ + # Create a tensor to store the gradients + grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] + if self.rnd: + grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None] + all_grads = torch.cat(grads) + + # Average the gradients across all GPUs + torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) + all_grads /= self.gpu_world_size + + # Get all parameters + all_params = self.policy.parameters() + if self.rnd: + all_params = chain(all_params, self.rnd.parameters()) + + # Update the gradients for all parameters with the reduced gradients + offset = 0 + for param in all_params: + if param.grad is not None: + numel = param.numel() + # copy data back from shared buffer + param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) + # update the offset for the next parameter + offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/__init__.py b/rsl_rl/build/lib/rsl_rl/__init__.py new file mode 100644 index 0000000..ebc2e20 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Main module for the rsl_rl package.""" diff --git a/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py b/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py new file mode 100644 index 0000000..0ef9163 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Implementation of different RL agents.""" + +from .distillation import Distillation +from .ppo import PPO +from .amp_ppo import AMP_PPO +__all__ = ["PPO", "Distillation", "AMP_PPO"] diff --git a/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py b/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py new file mode 100644 index 0000000..a2164eb --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py @@ -0,0 +1,571 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations +from torch._tensor import Tensor +from torch._tensor import Tensor +from typing import Any +from copy import deepcopy + +import torch +import torch.nn as nn +import torch.optim as optim +from itertools import chain + +from rsl_rl.modules import ActorCritic +from rsl_rl.modules.rnd import RandomNetworkDistillation +from rsl_rl.storage import RolloutStorage, ReplayBufferMulti +from rsl_rl.utils import string_to_callable + + +class AMP_PPO: + """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" + + policy: ActorCritic + """The actor critic module.""" + + def __init__( + self, + policy, + discriminator, + amp_data, + amp_normalizer, + amp_num_frames=1, + amp_replay_buffer_size=100000, + num_learning_epochs=5, + num_mini_batches=4, + clip_param=0.2, + gamma=0.99, + lam=0.95, + value_loss_coef=1.0, + entropy_coef=0.01, + learning_rate=0.001, + max_grad_norm=1.0, + use_clipped_value_loss=True, + schedule="adaptive", + desired_kl=0.01, + device="cpu", + normalize_advantage_per_mini_batch=False, + # RND parameters + rnd_cfg: dict | None = None, + # Symmetry parameters + symmetry_cfg: dict | None = None, + # Distributed training parameters + multi_gpu_cfg: dict | None = None, + ): + # device-related parameters + self.device = device + self.is_multi_gpu = multi_gpu_cfg is not None + # Multi-GPU parameters + if multi_gpu_cfg is not None: + self.gpu_global_rank = multi_gpu_cfg["global_rank"] + self.gpu_world_size = multi_gpu_cfg["world_size"] + else: + self.gpu_global_rank = 0 + self.gpu_world_size = 1 + + # RND components + if rnd_cfg is not None: + # Extract parameters used in ppo + rnd_lr = rnd_cfg.pop("learning_rate", 1e-3) + # Create RND module + self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) + # Create RND optimizer + params = self.rnd.predictor.parameters() + self.rnd_optimizer = optim.Adam(params, lr=rnd_lr) + else: + self.rnd = None + self.rnd_optimizer = None + + # Symmetry components + if symmetry_cfg is not None: + # Check if symmetry is enabled + use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] + # Print that we are not using symmetry + if not use_symmetry: + print("Symmetry not used for learning. We will use it for logging instead.") + # If function is a string then resolve it to a function + if isinstance(symmetry_cfg["data_augmentation_func"], str): + symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) + # Check valid configuration + if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): + raise ValueError( + "Data augmentation enabled but the function is not callable:" + f" {symmetry_cfg['data_augmentation_func']}" + ) + # Store symmetry configuration + self.symmetry = symmetry_cfg + else: + self.symmetry = None + + ## AMP components + self.discriminator = discriminator + self.discriminator.to(self.device) + + self.amp_storage = ReplayBufferMulti(discriminator.state_dim, amp_replay_buffer_size, amp_num_frames, device) + self.amp_data = amp_data + self.amp_normalizer = amp_normalizer + + # PPO components + self.policy = policy + self.policy.to(self.device) + + # Create rollout storage + self.storage: RolloutStorage = None # type: ignore + self.transition = RolloutStorage.Transition() + self.amp_transition = RolloutStorage.Transition() + params = [ + {'params': self.policy.parameters(), 'name': 'policy'}, + ] + + params.append({ + 'params': self.discriminator.trunk.parameters(), + 'weight_decay': 10e-4, + 'name': f'amp_trunk' + }) + params.append({ + 'params': self.discriminator.amp_linear.parameters(), + 'weight_decay': 10e-2, + 'name': f'amp_head' + }) + + # Create optimizer + self.optimizer = optim.Adam(params, lr=learning_rate) + + # PPO parameters + self.clip_param = clip_param + self.num_learning_epochs = num_learning_epochs + self.num_mini_batches = num_mini_batches + self.value_loss_coef = value_loss_coef + self.entropy_coef = entropy_coef + self.gamma = gamma + self.lam = lam + self.max_grad_norm = max_grad_norm + self.use_clipped_value_loss = use_clipped_value_loss + self.desired_kl = desired_kl + self.schedule = schedule + self.learning_rate = learning_rate + self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch + + + def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): + # create rollout storage + self.storage = RolloutStorage( + training_type, + num_envs, + num_transitions_per_env, + obs, + actions_shape, + self.device, + ) + + def act(self, obs, amp_obs): + if self.policy.is_recurrent: + self.transition.hidden_states = self.policy.get_hidden_states() + # compute the actions and values + self.transition.actions = self.policy.act(obs).detach() + self.transition.values = self.policy.evaluate(obs).detach() + self.transition.actions_log_prob = self.policy.get_actions_log_prob(self.transition.actions).detach() + self.transition.action_mean = self.policy.action_mean.detach() + self.transition.action_sigma = self.policy.action_std.detach() + # need to record obs before env.step() + self.transition.observations = obs + self.amp_transition.observations = amp_obs + return self.transition.actions + + def process_env_step(self, obs, rewards, dones, extras,amp_obs, amp_obs_frames=None): + # update the normalizers + self.policy.update_normalization(obs) + if self.rnd: + self.rnd.update_normalization(obs) + + # Record the rewards and dones + # Note: we clone here because later on we bootstrap the rewards based on timeouts + self.transition.rewards = rewards.clone() + self.transition.dones = dones + + # Compute the intrinsic rewards and add to extrinsic rewards + if self.rnd: + # Compute the intrinsic rewards + self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs) + # Add intrinsic rewards to extrinsic rewards + self.transition.rewards += self.intrinsic_rewards + + # Bootstrapping on time outs + if "time_outs" in extras: + self.transition.rewards += self.gamma * torch.squeeze( + self.transition.values * extras["time_outs"].unsqueeze(1).to(self.device), 1 + ) + + if amp_obs_frames is not None: + self.amp_storage.insert(amp_obs_frames) + else: + self.amp_storage.insert(self.amp_transition.observations, amp_obs) + + # record the transition + self.storage.add_transitions(self.transition) + self.transition.clear() + self.amp_transition.clear() + self.policy.reset(dones) + + def compute_returns(self, obs): + # compute value for the last step + last_values = self.policy.evaluate(obs).detach() + self.storage.compute_returns( + last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch + ) + + def update(self): # noqa: C901 + mean_value_loss = 0 + mean_surrogate_loss = 0 + mean_entropy = 0 + mean_amp_loss = 0 + mean_grad_pen_loss = 0 + mean_policy_pred = 0 + mean_expert_pred = 0 + # -- RND loss + if self.rnd: + mean_rnd_loss = 0 + else: + mean_rnd_loss = None + # -- Symmetry loss + if self.symmetry: + mean_symmetry_loss = 0 + else: + mean_symmetry_loss = None + + # generator for mini batches + if self.policy.is_recurrent: + generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + else: + generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + + + amp_policy_generator = self.amp_storage.feed_forward_generator( + self.num_learning_epochs * self.num_mini_batches, + self.storage.num_envs * self.storage.num_transitions_per_env // self.num_mini_batches, + ) + + amp_expert_generator = self.amp_data.feed_forward_generator_23dof_multi( + self.num_learning_epochs * self.num_mini_batches, + self.storage.num_envs * self.storage.num_transitions_per_env // self.num_mini_batches, + ) + + # iterate over batches + for sample, sample_amp_policy, sample_amp_expert in zip(generator, amp_policy_generator, amp_expert_generator): + ( + obs_batch, + actions_batch, + target_values_batch, + advantages_batch, + returns_batch, + old_actions_log_prob_batch, + old_mu_batch, + old_sigma_batch, + hid_states_batch, + masks_batch, + ) = sample + + # number of augmentations per sample + # we start with 1 and increase it if we use symmetry augmentation + num_aug = 1 + # original batch size + # we assume policy group is always there and needs augmentation + original_batch_size = obs_batch.batch_size[0] + + # check if we should normalize advantages per mini batch + if self.normalize_advantage_per_mini_batch: + with torch.no_grad(): + advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8) + + # Perform symmetric augmentation + if self.symmetry and self.symmetry["use_data_augmentation"]: + # augmentation using symmetry + data_augmentation_func = self.symmetry["data_augmentation_func"] + # returned shape: [batch_size * num_aug, ...] + obs_batch, actions_batch = data_augmentation_func( + obs=obs_batch, + actions=actions_batch, + env=self.symmetry["_env"], + ) + # compute number of augmentations per sample + # we assume policy group is always there and needs augmentation + num_aug = int(obs_batch.batch_size[0] / original_batch_size) + # repeat the rest of the batch + # -- actor + old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) + # -- critic + target_values_batch = target_values_batch.repeat(num_aug, 1) + advantages_batch = advantages_batch.repeat(num_aug, 1) + returns_batch = returns_batch.repeat(num_aug, 1) + + # Recompute actions log prob and entropy for current batch of transitions + # Note: we need to do this because we updated the policy with the new parameters + # -- actor + self.policy.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) + actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch) + # -- critic + value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) + # -- entropy + # we only keep the entropy of the first augmentation (the original one) + mu_batch = self.policy.action_mean[:original_batch_size] + sigma_batch = self.policy.action_std[:original_batch_size] + entropy_batch = self.policy.entropy[:original_batch_size] + + # KL + if self.desired_kl is not None and self.schedule == "adaptive": + with torch.inference_mode(): + kl = torch.sum( + torch.log(sigma_batch / old_sigma_batch + 1.0e-5) + + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) + / (2.0 * torch.square(sigma_batch)) + - 0.5, + axis=-1, + ) + kl_mean = torch.mean(kl) + + # Reduce the KL divergence across all GPUs + if self.is_multi_gpu: + torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM) + kl_mean /= self.gpu_world_size + + # Update the learning rate + # Perform this adaptation only on the main process + # TODO: Is this needed? If KL-divergence is the "same" across all GPUs, + # then the learning rate should be the same across all GPUs. + if self.gpu_global_rank == 0: + if kl_mean > self.desired_kl * 2.0: + self.learning_rate = max(1e-5, self.learning_rate / 1.5) + elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: + self.learning_rate = min(1e-2, self.learning_rate * 1.5) + + # Update the learning rate for all GPUs + if self.is_multi_gpu: + lr_tensor = torch.tensor(self.learning_rate, device=self.device) + torch.distributed.broadcast(lr_tensor, src=0) + self.learning_rate = lr_tensor.item() + + # Update the learning rate for all parameter groups + for param_group in self.optimizer.param_groups: + param_group["lr"] = self.learning_rate + + # Surrogate loss + ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) + surrogate = -torch.squeeze(advantages_batch) * ratio + surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( + ratio, 1.0 - self.clip_param, 1.0 + self.clip_param + ) + surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() + + # Value function loss + if self.use_clipped_value_loss: + value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( + -self.clip_param, self.clip_param + ) + value_losses = (value_batch - returns_batch).pow(2) + value_losses_clipped = (value_clipped - returns_batch).pow(2) + value_loss = torch.max(value_losses, value_losses_clipped).mean() + else: + value_loss = (returns_batch - value_batch).pow(2).mean() + + loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() + + # Symmetry loss + if self.symmetry: + # obtain the symmetric actions + # if we did augmentation before then we don't need to augment again + if not self.symmetry["use_data_augmentation"]: + data_augmentation_func = self.symmetry["data_augmentation_func"] + obs_batch, _ = data_augmentation_func(obs=obs_batch, actions=None, env=self.symmetry["_env"]) + # compute number of augmentations per sample + num_aug = int(obs_batch.shape[0] / original_batch_size) + + # actions predicted by the actor for symmetrically-augmented observations + mean_actions_batch = self.policy.act_inference(obs_batch.detach().clone()) + + # compute the symmetrically augmented actions + # note: we are assuming the first augmentation is the original one. + # We do not use the action_batch from earlier since that action was sampled from the distribution. + # However, the symmetry loss is computed using the mean of the distribution. + action_mean_orig = mean_actions_batch[:original_batch_size] + _, actions_mean_symm_batch = data_augmentation_func( + obs=None, actions=action_mean_orig, env=self.symmetry["_env"] + ) + + # compute the loss (we skip the first augmentation as it is the original one) + mse_loss = torch.nn.MSELoss() + symmetry_loss = mse_loss( + mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] + ) + # add the loss to the total loss + if self.symmetry["use_mirror_loss"]: + loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss + else: + symmetry_loss = symmetry_loss.detach() + + # Random Network Distillation loss + # TODO: Move this processing to inside RND module. + if self.rnd: + # extract the rnd_state + # TODO: Check if we still need torch no grad. It is just an affine transformation. + with torch.no_grad(): + rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size]) + rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch) + # predict the embedding and the target + predicted_embedding = self.rnd.predictor(rnd_state_batch) + target_embedding = self.rnd.target(rnd_state_batch).detach() + # compute the loss as the mean squared error + mseloss = torch.nn.MSELoss() + rnd_loss = mseloss(predicted_embedding, target_embedding) + + expert_states = sample_amp_expert + policy_states = sample_amp_policy + + with torch.no_grad(): + expert_states = self.amp_normalizer.normalize_torch(expert_states.to(self.device), self.device) + policy_states = self.amp_normalizer.normalize_torch(policy_states, self.device) + + contact_phase_push = obs_batch['critic'][:, -4] + mask_push = contact_phase_push == 1. + + if mask_push.any(): + policy_d = self.discriminator(policy_states.flatten(1)) + expert_states = expert_states.to(self.device) + expert_d = self.discriminator(expert_states.flatten(1)) + + expert_loss = torch.nn.MSELoss()(expert_d, torch.ones(expert_d.size(), device=self.device)) + policy_loss = torch.nn.MSELoss()(policy_d, -1 * torch.ones(policy_d.size(), device=self.device)) + amp_loss = 0.5 * (expert_loss + policy_loss) + + # grad penalty + grad_pen_loss = self.discriminator.compute_grad_pen(expert_states, lambda_=5) + else: + amp_loss = torch.tensor(0.0, device=self.device) + grad_pen_loss = torch.tensor(0.0, device=self.device) + expert_loss = torch.tensor(0.0, device=self.device) + policy_loss = torch.tensor(0.0, device=self.device) + + loss += (amp_loss + grad_pen_loss) + self.amp_normalizer.update(policy_states.cpu().numpy()) + self.amp_normalizer.update(expert_states.cpu().numpy()) + + # Compute the gradients + # -- For PPO + self.optimizer.zero_grad() + loss.backward() + # -- For RND + if self.rnd: + self.rnd_optimizer.zero_grad() # type: ignore + rnd_loss.backward() + + # Collect gradients from all GPUs + if self.is_multi_gpu: + self.reduce_parameters() + + # Apply the gradients + # -- For PPO + nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) + self.optimizer.step() + # -- For RND + if self.rnd_optimizer: + self.rnd_optimizer.step() + + # Store the losses + mean_value_loss += value_loss.item() + mean_surrogate_loss += surrogate_loss.item() + mean_entropy += entropy_batch.mean().item() + mean_amp_loss += amp_loss.item() + mean_grad_pen_loss += grad_pen_loss.item() + mean_policy_pred += policy_loss.mean().item() + mean_expert_pred += expert_loss.mean().item() + # -- RND loss + if mean_rnd_loss is not None: + mean_rnd_loss += rnd_loss.item() + # -- Symmetry loss + if mean_symmetry_loss is not None: + mean_symmetry_loss += symmetry_loss.item() + + # -- For PPO + num_updates = self.num_learning_epochs * self.num_mini_batches + mean_value_loss /= num_updates + mean_surrogate_loss /= num_updates + mean_entropy /= num_updates + mean_amp_loss /= num_updates + mean_grad_pen_loss /= num_updates + mean_policy_pred /= num_updates + mean_expert_pred /= num_updates + + # -- For RND + if mean_rnd_loss is not None: + mean_rnd_loss /= num_updates + # -- For Symmetry + if mean_symmetry_loss is not None: + mean_symmetry_loss /= num_updates + # -- Clear the storage + self.storage.clear() + + # construct the loss dictionary + loss_dict = { + "value_function": mean_value_loss, + "surrogate": mean_surrogate_loss, + "entropy": mean_entropy, + "amp": mean_amp_loss, + "amp_grad_pen": mean_grad_pen_loss, + "amp_policy_pred": mean_policy_pred, + "amp_expert_pred": mean_expert_pred, + } + if self.rnd: + loss_dict["rnd"] = mean_rnd_loss + if self.symmetry: + loss_dict["symmetry"] = mean_symmetry_loss + + return loss_dict + + """ + Helper functions + """ + + def broadcast_parameters(self): + """Broadcast model parameters to all GPUs.""" + # obtain the model parameters on current GPU + model_params = [self.policy.state_dict()] + if self.rnd: + model_params.append(self.rnd.predictor.state_dict()) + # broadcast the model parameters + torch.distributed.broadcast_object_list(model_params, src=0) + # load the model parameters on all GPUs from source GPU + self.policy.load_state_dict(model_params[0]) + if self.rnd: + self.rnd.predictor.load_state_dict(model_params[1]) + + def reduce_parameters(self): + """Collect gradients from all GPUs and average them. + + This function is called after the backward pass to synchronize the gradients across all GPUs. + """ + # Create a tensor to store the gradients + grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] + if self.rnd: + grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None] + all_grads = torch.cat(grads) + + # Average the gradients across all GPUs + torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) + all_grads /= self.gpu_world_size + + # Get all parameters + all_params = self.policy.parameters() + if self.rnd: + all_params = chain(all_params, self.rnd.parameters()) + + # Update the gradients for all parameters with the reduced gradients + offset = 0 + for param in all_params: + if param.grad is not None: + numel = param.numel() + # copy data back from shared buffer + param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) + # update the offset for the next parameter + offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py b/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py new file mode 100644 index 0000000..3a86e00 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py @@ -0,0 +1,185 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch +import torch.nn as nn + +from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent +from rsl_rl.storage import RolloutStorage +from rsl_rl.utils import resolve_optimizer + + +class Distillation: + """Distillation algorithm for training a student model to mimic a teacher model.""" + + policy: StudentTeacher | StudentTeacherRecurrent + """The student teacher model.""" + + def __init__( + self, + policy, + num_learning_epochs=1, + gradient_length=15, + learning_rate=1e-3, + max_grad_norm=None, + loss_type="mse", + optimizer="adam", + device="cpu", + # Distributed training parameters + multi_gpu_cfg: dict | None = None, + ): + # device-related parameters + self.device = device + self.is_multi_gpu = multi_gpu_cfg is not None + # Multi-GPU parameters + if multi_gpu_cfg is not None: + self.gpu_global_rank = multi_gpu_cfg["global_rank"] + self.gpu_world_size = multi_gpu_cfg["world_size"] + else: + self.gpu_global_rank = 0 + self.gpu_world_size = 1 + + # distillation components + self.policy = policy + self.policy.to(self.device) + self.storage = None # initialized later + + # initialize the optimizer + self.optimizer = resolve_optimizer(optimizer)(self.policy.parameters(), lr=learning_rate) + + # initialize the transition + self.transition = RolloutStorage.Transition() + self.last_hidden_states = None + + # distillation parameters + self.num_learning_epochs = num_learning_epochs + self.gradient_length = gradient_length + self.learning_rate = learning_rate + self.max_grad_norm = max_grad_norm + + # initialize the loss function + loss_fn_dict = { + "mse": nn.functional.mse_loss, + "huber": nn.functional.huber_loss, + } + if loss_type in loss_fn_dict: + self.loss_fn = loss_fn_dict[loss_type] + else: + raise ValueError(f"Unknown loss type: {loss_type}. Supported types are: {list(loss_fn_dict.keys())}") + + self.num_updates = 0 + + def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): + # create rollout storage + self.storage = RolloutStorage( + training_type, + num_envs, + num_transitions_per_env, + obs, + actions_shape, + self.device, + ) + + def act(self, obs): + # compute the actions + self.transition.actions = self.policy.act(obs).detach() + self.transition.privileged_actions = self.policy.evaluate(obs).detach() + # record the observations + self.transition.observations = obs + return self.transition.actions + + def process_env_step(self, obs, rewards, dones, extras): + # update the normalizers + self.policy.update_normalization(obs) + + # record the rewards and dones + self.transition.rewards = rewards + self.transition.dones = dones + # record the transition + self.storage.add_transitions(self.transition) + self.transition.clear() + self.policy.reset(dones) + + def update(self): + self.num_updates += 1 + mean_behavior_loss = 0 + loss = 0 + cnt = 0 + + for epoch in range(self.num_learning_epochs): + self.policy.reset(hidden_states=self.last_hidden_states) + self.policy.detach_hidden_states() + for obs, _, privileged_actions, dones in self.storage.generator(): + + # inference the student for gradient computation + actions = self.policy.act_inference(obs) + + # behavior cloning loss + behavior_loss = self.loss_fn(actions, privileged_actions) + + # total loss + loss = loss + behavior_loss + mean_behavior_loss += behavior_loss.item() + cnt += 1 + + # gradient step + if cnt % self.gradient_length == 0: + self.optimizer.zero_grad() + loss.backward() + if self.is_multi_gpu: + self.reduce_parameters() + if self.max_grad_norm: + nn.utils.clip_grad_norm_(self.policy.student.parameters(), self.max_grad_norm) + self.optimizer.step() + self.policy.detach_hidden_states() + loss = 0 + + # reset dones + self.policy.reset(dones.view(-1)) + self.policy.detach_hidden_states(dones.view(-1)) + + mean_behavior_loss /= cnt + self.storage.clear() + self.last_hidden_states = self.policy.get_hidden_states() + self.policy.detach_hidden_states() + + # construct the loss dictionary + loss_dict = {"behavior": mean_behavior_loss} + + return loss_dict + + """ + Helper functions + """ + + def broadcast_parameters(self): + """Broadcast model parameters to all GPUs.""" + # obtain the model parameters on current GPU + model_params = [self.policy.state_dict()] + # broadcast the model parameters + torch.distributed.broadcast_object_list(model_params, src=0) + # load the model parameters on all GPUs from source GPU + self.policy.load_state_dict(model_params[0]) + + def reduce_parameters(self): + """Collect gradients from all GPUs and average them. + + This function is called after the backward pass to synchronize the gradients across all GPUs. + """ + # Create a tensor to store the gradients + grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] + all_grads = torch.cat(grads) + # Average the gradients across all GPUs + torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) + all_grads /= self.gpu_world_size + # Update the gradients for all parameters with the reduced gradients + offset = 0 + for param in self.policy.parameters(): + if param.grad is not None: + numel = param.numel() + # copy data back from shared buffer + param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) + # update the offset for the next parameter + offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py b/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py new file mode 100644 index 0000000..6c21fc5 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py @@ -0,0 +1,469 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.optim as optim +from itertools import chain + +from rsl_rl.modules import ActorCritic +from rsl_rl.modules.rnd import RandomNetworkDistillation +from rsl_rl.storage import RolloutStorage +from rsl_rl.utils import string_to_callable + + +class PPO: + """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" + + policy: ActorCritic + """The actor critic module.""" + + def __init__( + self, + policy, + num_learning_epochs=5, + num_mini_batches=4, + clip_param=0.2, + gamma=0.99, + lam=0.95, + value_loss_coef=1.0, + entropy_coef=0.01, + learning_rate=0.001, + max_grad_norm=1.0, + use_clipped_value_loss=True, + schedule="adaptive", + desired_kl=0.01, + device="cpu", + normalize_advantage_per_mini_batch=False, + # RND parameters + rnd_cfg: dict | None = None, + # Symmetry parameters + symmetry_cfg: dict | None = None, + # Distributed training parameters + multi_gpu_cfg: dict | None = None, + ): + # device-related parameters + self.device = device + self.is_multi_gpu = multi_gpu_cfg is not None + # Multi-GPU parameters + if multi_gpu_cfg is not None: + self.gpu_global_rank = multi_gpu_cfg["global_rank"] + self.gpu_world_size = multi_gpu_cfg["world_size"] + else: + self.gpu_global_rank = 0 + self.gpu_world_size = 1 + + # RND components + if rnd_cfg is not None: + # Extract parameters used in ppo + rnd_lr = rnd_cfg.pop("learning_rate", 1e-3) + # Create RND module + self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) + # Create RND optimizer + params = self.rnd.predictor.parameters() + self.rnd_optimizer = optim.Adam(params, lr=rnd_lr) + else: + self.rnd = None + self.rnd_optimizer = None + + # Symmetry components + if symmetry_cfg is not None: + # Check if symmetry is enabled + use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] + # Print that we are not using symmetry + if not use_symmetry: + print("Symmetry not used for learning. We will use it for logging instead.") + # If function is a string then resolve it to a function + if isinstance(symmetry_cfg["data_augmentation_func"], str): + symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) + # Check valid configuration + if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): + raise ValueError( + "Data augmentation enabled but the function is not callable:" + f" {symmetry_cfg['data_augmentation_func']}" + ) + # Store symmetry configuration + self.symmetry = symmetry_cfg + else: + self.symmetry = None + + # PPO components + self.policy = policy + self.policy.to(self.device) + # Create optimizer + self.optimizer = optim.Adam(self.policy.parameters(), lr=learning_rate) + # Create rollout storage + self.storage: RolloutStorage = None # type: ignore + self.transition = RolloutStorage.Transition() + + # PPO parameters + self.clip_param = clip_param + self.num_learning_epochs = num_learning_epochs + self.num_mini_batches = num_mini_batches + self.value_loss_coef = value_loss_coef + self.entropy_coef = entropy_coef + self.gamma = gamma + self.lam = lam + self.max_grad_norm = max_grad_norm + self.use_clipped_value_loss = use_clipped_value_loss + self.desired_kl = desired_kl + self.schedule = schedule + self.learning_rate = learning_rate + self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch + + def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): + # create rollout storage + self.storage = RolloutStorage( + training_type, + num_envs, + num_transitions_per_env, + obs, + actions_shape, + self.device, + ) + + def act(self, obs): + if self.policy.is_recurrent: + self.transition.hidden_states = self.policy.get_hidden_states() + # compute the actions and values + self.transition.actions = self.policy.act(obs).detach() + self.transition.values = self.policy.evaluate(obs).detach() + self.transition.actions_log_prob = self.policy.get_actions_log_prob(self.transition.actions).detach() + self.transition.action_mean = self.policy.action_mean.detach() + self.transition.action_sigma = self.policy.action_std.detach() + # need to record obs before env.step() + self.transition.observations = obs + return self.transition.actions + + def process_env_step(self, obs, rewards, dones, extras): + # update the normalizers + self.policy.update_normalization(obs) + if self.rnd: + self.rnd.update_normalization(obs) + + # Record the rewards and dones + # Note: we clone here because later on we bootstrap the rewards based on timeouts + self.transition.rewards = rewards.clone() + self.transition.dones = dones + + # Compute the intrinsic rewards and add to extrinsic rewards + if self.rnd: + # Compute the intrinsic rewards + self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs) + # Add intrinsic rewards to extrinsic rewards + self.transition.rewards += self.intrinsic_rewards + + # Bootstrapping on time outs + if "time_outs" in extras: + self.transition.rewards += self.gamma * torch.squeeze( + self.transition.values * extras["time_outs"].unsqueeze(1).to(self.device), 1 + ) + + # record the transition + self.storage.add_transitions(self.transition) + self.transition.clear() + self.policy.reset(dones) + + def compute_returns(self, obs): + # compute value for the last step + last_values = self.policy.evaluate(obs).detach() + self.storage.compute_returns( + last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch + ) + + def update(self): # noqa: C901 + mean_value_loss = 0 + mean_surrogate_loss = 0 + mean_entropy = 0 + # -- RND loss + if self.rnd: + mean_rnd_loss = 0 + else: + mean_rnd_loss = None + # -- Symmetry loss + if self.symmetry: + mean_symmetry_loss = 0 + else: + mean_symmetry_loss = None + + # generator for mini batches + if self.policy.is_recurrent: + generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + else: + generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + + # iterate over batches + for ( + obs_batch, + actions_batch, + target_values_batch, + advantages_batch, + returns_batch, + old_actions_log_prob_batch, + old_mu_batch, + old_sigma_batch, + hid_states_batch, + masks_batch, + ) in generator: + + # number of augmentations per sample + # we start with 1 and increase it if we use symmetry augmentation + num_aug = 1 + # original batch size + # we assume policy group is always there and needs augmentation + original_batch_size = obs_batch.batch_size[0] + + # check if we should normalize advantages per mini batch + if self.normalize_advantage_per_mini_batch: + with torch.no_grad(): + advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8) + + # Perform symmetric augmentation + if self.symmetry and self.symmetry["use_data_augmentation"]: + # augmentation using symmetry + data_augmentation_func = self.symmetry["data_augmentation_func"] + # returned shape: [batch_size * num_aug, ...] + obs_batch, actions_batch = data_augmentation_func( + obs=obs_batch, + actions=actions_batch, + env=self.symmetry["_env"], + ) + # compute number of augmentations per sample + # we assume policy group is always there and needs augmentation + num_aug = int(obs_batch.batch_size[0] / original_batch_size) + # repeat the rest of the batch + # -- actor + old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) + # -- critic + target_values_batch = target_values_batch.repeat(num_aug, 1) + advantages_batch = advantages_batch.repeat(num_aug, 1) + returns_batch = returns_batch.repeat(num_aug, 1) + + # Recompute actions log prob and entropy for current batch of transitions + # Note: we need to do this because we updated the policy with the new parameters + # -- actor + self.policy.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) + actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch) + # -- critic + value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) + # -- entropy + # we only keep the entropy of the first augmentation (the original one) + mu_batch = self.policy.action_mean[:original_batch_size] + sigma_batch = self.policy.action_std[:original_batch_size] + entropy_batch = self.policy.entropy[:original_batch_size] + + # KL + if self.desired_kl is not None and self.schedule == "adaptive": + with torch.inference_mode(): + kl = torch.sum( + torch.log(sigma_batch / old_sigma_batch + 1.0e-5) + + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) + / (2.0 * torch.square(sigma_batch)) + - 0.5, + axis=-1, + ) + kl_mean = torch.mean(kl) + + # Reduce the KL divergence across all GPUs + if self.is_multi_gpu: + torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM) + kl_mean /= self.gpu_world_size + + # Update the learning rate + # Perform this adaptation only on the main process + # TODO: Is this needed? If KL-divergence is the "same" across all GPUs, + # then the learning rate should be the same across all GPUs. + if self.gpu_global_rank == 0: + if kl_mean > self.desired_kl * 2.0: + self.learning_rate = max(1e-5, self.learning_rate / 1.5) + elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: + self.learning_rate = min(1e-2, self.learning_rate * 1.5) + + # Update the learning rate for all GPUs + if self.is_multi_gpu: + lr_tensor = torch.tensor(self.learning_rate, device=self.device) + torch.distributed.broadcast(lr_tensor, src=0) + self.learning_rate = lr_tensor.item() + + # Update the learning rate for all parameter groups + for param_group in self.optimizer.param_groups: + param_group["lr"] = self.learning_rate + + # Surrogate loss + ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) + surrogate = -torch.squeeze(advantages_batch) * ratio + surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( + ratio, 1.0 - self.clip_param, 1.0 + self.clip_param + ) + surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() + + # Value function loss + if self.use_clipped_value_loss: + value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( + -self.clip_param, self.clip_param + ) + value_losses = (value_batch - returns_batch).pow(2) + value_losses_clipped = (value_clipped - returns_batch).pow(2) + value_loss = torch.max(value_losses, value_losses_clipped).mean() + else: + value_loss = (returns_batch - value_batch).pow(2).mean() + + loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() + + # Symmetry loss + if self.symmetry: + # obtain the symmetric actions + # if we did augmentation before then we don't need to augment again + if not self.symmetry["use_data_augmentation"]: + data_augmentation_func = self.symmetry["data_augmentation_func"] + obs_batch, _ = data_augmentation_func(obs=obs_batch, actions=None, env=self.symmetry["_env"]) + # compute number of augmentations per sample + num_aug = int(obs_batch.shape[0] / original_batch_size) + + # actions predicted by the actor for symmetrically-augmented observations + mean_actions_batch = self.policy.act_inference(obs_batch.detach().clone()) + + # compute the symmetrically augmented actions + # note: we are assuming the first augmentation is the original one. + # We do not use the action_batch from earlier since that action was sampled from the distribution. + # However, the symmetry loss is computed using the mean of the distribution. + action_mean_orig = mean_actions_batch[:original_batch_size] + _, actions_mean_symm_batch = data_augmentation_func( + obs=None, actions=action_mean_orig, env=self.symmetry["_env"] + ) + + # compute the loss (we skip the first augmentation as it is the original one) + mse_loss = torch.nn.MSELoss() + symmetry_loss = mse_loss( + mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] + ) + # add the loss to the total loss + if self.symmetry["use_mirror_loss"]: + loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss + else: + symmetry_loss = symmetry_loss.detach() + + # Random Network Distillation loss + # TODO: Move this processing to inside RND module. + if self.rnd: + # extract the rnd_state + # TODO: Check if we still need torch no grad. It is just an affine transformation. + with torch.no_grad(): + rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size]) + rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch) + # predict the embedding and the target + predicted_embedding = self.rnd.predictor(rnd_state_batch) + target_embedding = self.rnd.target(rnd_state_batch).detach() + # compute the loss as the mean squared error + mseloss = torch.nn.MSELoss() + rnd_loss = mseloss(predicted_embedding, target_embedding) + + # Compute the gradients + # -- For PPO + self.optimizer.zero_grad() + loss.backward() + # -- For RND + if self.rnd: + self.rnd_optimizer.zero_grad() # type: ignore + rnd_loss.backward() + + # Collect gradients from all GPUs + if self.is_multi_gpu: + self.reduce_parameters() + + # Apply the gradients + # -- For PPO + nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) + self.optimizer.step() + # -- For RND + if self.rnd_optimizer: + self.rnd_optimizer.step() + + # Store the losses + mean_value_loss += value_loss.item() + mean_surrogate_loss += surrogate_loss.item() + mean_entropy += entropy_batch.mean().item() + # -- RND loss + if mean_rnd_loss is not None: + mean_rnd_loss += rnd_loss.item() + # -- Symmetry loss + if mean_symmetry_loss is not None: + mean_symmetry_loss += symmetry_loss.item() + + # -- For PPO + num_updates = self.num_learning_epochs * self.num_mini_batches + mean_value_loss /= num_updates + mean_surrogate_loss /= num_updates + mean_entropy /= num_updates + # -- For RND + if mean_rnd_loss is not None: + mean_rnd_loss /= num_updates + # -- For Symmetry + if mean_symmetry_loss is not None: + mean_symmetry_loss /= num_updates + # -- Clear the storage + self.storage.clear() + + # construct the loss dictionary + loss_dict = { + "value_function": mean_value_loss, + "surrogate": mean_surrogate_loss, + "entropy": mean_entropy, + } + if self.rnd: + loss_dict["rnd"] = mean_rnd_loss + if self.symmetry: + loss_dict["symmetry"] = mean_symmetry_loss + + return loss_dict + + """ + Helper functions + """ + + def broadcast_parameters(self): + """Broadcast model parameters to all GPUs.""" + # obtain the model parameters on current GPU + model_params = [self.policy.state_dict()] + if self.rnd: + model_params.append(self.rnd.predictor.state_dict()) + # broadcast the model parameters + torch.distributed.broadcast_object_list(model_params, src=0) + # load the model parameters on all GPUs from source GPU + self.policy.load_state_dict(model_params[0]) + if self.rnd: + self.rnd.predictor.load_state_dict(model_params[1]) + + def reduce_parameters(self): + """Collect gradients from all GPUs and average them. + + This function is called after the backward pass to synchronize the gradients across all GPUs. + """ + # Create a tensor to store the gradients + grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] + if self.rnd: + grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None] + all_grads = torch.cat(grads) + + # Average the gradients across all GPUs + torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) + all_grads /= self.gpu_world_size + + # Get all parameters + all_params = self.policy.parameters() + if self.rnd: + all_params = chain(all_params, self.rnd.parameters()) + + # Update the gradients for all parameters with the reduced gradients + offset = 0 + for param in all_params: + if param.grad is not None: + numel = param.numel() + # copy data back from shared buffer + param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) + # update the offset for the next parameter + offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/__init__.py new file mode 100644 index 0000000..ebc2e20 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Main module for the rsl_rl package.""" diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py new file mode 100644 index 0000000..0ef9163 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Implementation of different RL agents.""" + +from .distillation import Distillation +from .ppo import PPO +from .amp_ppo import AMP_PPO +__all__ = ["PPO", "Distillation", "AMP_PPO"] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py new file mode 100644 index 0000000..a2164eb --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py @@ -0,0 +1,571 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations +from torch._tensor import Tensor +from torch._tensor import Tensor +from typing import Any +from copy import deepcopy + +import torch +import torch.nn as nn +import torch.optim as optim +from itertools import chain + +from rsl_rl.modules import ActorCritic +from rsl_rl.modules.rnd import RandomNetworkDistillation +from rsl_rl.storage import RolloutStorage, ReplayBufferMulti +from rsl_rl.utils import string_to_callable + + +class AMP_PPO: + """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" + + policy: ActorCritic + """The actor critic module.""" + + def __init__( + self, + policy, + discriminator, + amp_data, + amp_normalizer, + amp_num_frames=1, + amp_replay_buffer_size=100000, + num_learning_epochs=5, + num_mini_batches=4, + clip_param=0.2, + gamma=0.99, + lam=0.95, + value_loss_coef=1.0, + entropy_coef=0.01, + learning_rate=0.001, + max_grad_norm=1.0, + use_clipped_value_loss=True, + schedule="adaptive", + desired_kl=0.01, + device="cpu", + normalize_advantage_per_mini_batch=False, + # RND parameters + rnd_cfg: dict | None = None, + # Symmetry parameters + symmetry_cfg: dict | None = None, + # Distributed training parameters + multi_gpu_cfg: dict | None = None, + ): + # device-related parameters + self.device = device + self.is_multi_gpu = multi_gpu_cfg is not None + # Multi-GPU parameters + if multi_gpu_cfg is not None: + self.gpu_global_rank = multi_gpu_cfg["global_rank"] + self.gpu_world_size = multi_gpu_cfg["world_size"] + else: + self.gpu_global_rank = 0 + self.gpu_world_size = 1 + + # RND components + if rnd_cfg is not None: + # Extract parameters used in ppo + rnd_lr = rnd_cfg.pop("learning_rate", 1e-3) + # Create RND module + self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) + # Create RND optimizer + params = self.rnd.predictor.parameters() + self.rnd_optimizer = optim.Adam(params, lr=rnd_lr) + else: + self.rnd = None + self.rnd_optimizer = None + + # Symmetry components + if symmetry_cfg is not None: + # Check if symmetry is enabled + use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] + # Print that we are not using symmetry + if not use_symmetry: + print("Symmetry not used for learning. We will use it for logging instead.") + # If function is a string then resolve it to a function + if isinstance(symmetry_cfg["data_augmentation_func"], str): + symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) + # Check valid configuration + if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): + raise ValueError( + "Data augmentation enabled but the function is not callable:" + f" {symmetry_cfg['data_augmentation_func']}" + ) + # Store symmetry configuration + self.symmetry = symmetry_cfg + else: + self.symmetry = None + + ## AMP components + self.discriminator = discriminator + self.discriminator.to(self.device) + + self.amp_storage = ReplayBufferMulti(discriminator.state_dim, amp_replay_buffer_size, amp_num_frames, device) + self.amp_data = amp_data + self.amp_normalizer = amp_normalizer + + # PPO components + self.policy = policy + self.policy.to(self.device) + + # Create rollout storage + self.storage: RolloutStorage = None # type: ignore + self.transition = RolloutStorage.Transition() + self.amp_transition = RolloutStorage.Transition() + params = [ + {'params': self.policy.parameters(), 'name': 'policy'}, + ] + + params.append({ + 'params': self.discriminator.trunk.parameters(), + 'weight_decay': 10e-4, + 'name': f'amp_trunk' + }) + params.append({ + 'params': self.discriminator.amp_linear.parameters(), + 'weight_decay': 10e-2, + 'name': f'amp_head' + }) + + # Create optimizer + self.optimizer = optim.Adam(params, lr=learning_rate) + + # PPO parameters + self.clip_param = clip_param + self.num_learning_epochs = num_learning_epochs + self.num_mini_batches = num_mini_batches + self.value_loss_coef = value_loss_coef + self.entropy_coef = entropy_coef + self.gamma = gamma + self.lam = lam + self.max_grad_norm = max_grad_norm + self.use_clipped_value_loss = use_clipped_value_loss + self.desired_kl = desired_kl + self.schedule = schedule + self.learning_rate = learning_rate + self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch + + + def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): + # create rollout storage + self.storage = RolloutStorage( + training_type, + num_envs, + num_transitions_per_env, + obs, + actions_shape, + self.device, + ) + + def act(self, obs, amp_obs): + if self.policy.is_recurrent: + self.transition.hidden_states = self.policy.get_hidden_states() + # compute the actions and values + self.transition.actions = self.policy.act(obs).detach() + self.transition.values = self.policy.evaluate(obs).detach() + self.transition.actions_log_prob = self.policy.get_actions_log_prob(self.transition.actions).detach() + self.transition.action_mean = self.policy.action_mean.detach() + self.transition.action_sigma = self.policy.action_std.detach() + # need to record obs before env.step() + self.transition.observations = obs + self.amp_transition.observations = amp_obs + return self.transition.actions + + def process_env_step(self, obs, rewards, dones, extras,amp_obs, amp_obs_frames=None): + # update the normalizers + self.policy.update_normalization(obs) + if self.rnd: + self.rnd.update_normalization(obs) + + # Record the rewards and dones + # Note: we clone here because later on we bootstrap the rewards based on timeouts + self.transition.rewards = rewards.clone() + self.transition.dones = dones + + # Compute the intrinsic rewards and add to extrinsic rewards + if self.rnd: + # Compute the intrinsic rewards + self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs) + # Add intrinsic rewards to extrinsic rewards + self.transition.rewards += self.intrinsic_rewards + + # Bootstrapping on time outs + if "time_outs" in extras: + self.transition.rewards += self.gamma * torch.squeeze( + self.transition.values * extras["time_outs"].unsqueeze(1).to(self.device), 1 + ) + + if amp_obs_frames is not None: + self.amp_storage.insert(amp_obs_frames) + else: + self.amp_storage.insert(self.amp_transition.observations, amp_obs) + + # record the transition + self.storage.add_transitions(self.transition) + self.transition.clear() + self.amp_transition.clear() + self.policy.reset(dones) + + def compute_returns(self, obs): + # compute value for the last step + last_values = self.policy.evaluate(obs).detach() + self.storage.compute_returns( + last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch + ) + + def update(self): # noqa: C901 + mean_value_loss = 0 + mean_surrogate_loss = 0 + mean_entropy = 0 + mean_amp_loss = 0 + mean_grad_pen_loss = 0 + mean_policy_pred = 0 + mean_expert_pred = 0 + # -- RND loss + if self.rnd: + mean_rnd_loss = 0 + else: + mean_rnd_loss = None + # -- Symmetry loss + if self.symmetry: + mean_symmetry_loss = 0 + else: + mean_symmetry_loss = None + + # generator for mini batches + if self.policy.is_recurrent: + generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + else: + generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + + + amp_policy_generator = self.amp_storage.feed_forward_generator( + self.num_learning_epochs * self.num_mini_batches, + self.storage.num_envs * self.storage.num_transitions_per_env // self.num_mini_batches, + ) + + amp_expert_generator = self.amp_data.feed_forward_generator_23dof_multi( + self.num_learning_epochs * self.num_mini_batches, + self.storage.num_envs * self.storage.num_transitions_per_env // self.num_mini_batches, + ) + + # iterate over batches + for sample, sample_amp_policy, sample_amp_expert in zip(generator, amp_policy_generator, amp_expert_generator): + ( + obs_batch, + actions_batch, + target_values_batch, + advantages_batch, + returns_batch, + old_actions_log_prob_batch, + old_mu_batch, + old_sigma_batch, + hid_states_batch, + masks_batch, + ) = sample + + # number of augmentations per sample + # we start with 1 and increase it if we use symmetry augmentation + num_aug = 1 + # original batch size + # we assume policy group is always there and needs augmentation + original_batch_size = obs_batch.batch_size[0] + + # check if we should normalize advantages per mini batch + if self.normalize_advantage_per_mini_batch: + with torch.no_grad(): + advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8) + + # Perform symmetric augmentation + if self.symmetry and self.symmetry["use_data_augmentation"]: + # augmentation using symmetry + data_augmentation_func = self.symmetry["data_augmentation_func"] + # returned shape: [batch_size * num_aug, ...] + obs_batch, actions_batch = data_augmentation_func( + obs=obs_batch, + actions=actions_batch, + env=self.symmetry["_env"], + ) + # compute number of augmentations per sample + # we assume policy group is always there and needs augmentation + num_aug = int(obs_batch.batch_size[0] / original_batch_size) + # repeat the rest of the batch + # -- actor + old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) + # -- critic + target_values_batch = target_values_batch.repeat(num_aug, 1) + advantages_batch = advantages_batch.repeat(num_aug, 1) + returns_batch = returns_batch.repeat(num_aug, 1) + + # Recompute actions log prob and entropy for current batch of transitions + # Note: we need to do this because we updated the policy with the new parameters + # -- actor + self.policy.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) + actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch) + # -- critic + value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) + # -- entropy + # we only keep the entropy of the first augmentation (the original one) + mu_batch = self.policy.action_mean[:original_batch_size] + sigma_batch = self.policy.action_std[:original_batch_size] + entropy_batch = self.policy.entropy[:original_batch_size] + + # KL + if self.desired_kl is not None and self.schedule == "adaptive": + with torch.inference_mode(): + kl = torch.sum( + torch.log(sigma_batch / old_sigma_batch + 1.0e-5) + + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) + / (2.0 * torch.square(sigma_batch)) + - 0.5, + axis=-1, + ) + kl_mean = torch.mean(kl) + + # Reduce the KL divergence across all GPUs + if self.is_multi_gpu: + torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM) + kl_mean /= self.gpu_world_size + + # Update the learning rate + # Perform this adaptation only on the main process + # TODO: Is this needed? If KL-divergence is the "same" across all GPUs, + # then the learning rate should be the same across all GPUs. + if self.gpu_global_rank == 0: + if kl_mean > self.desired_kl * 2.0: + self.learning_rate = max(1e-5, self.learning_rate / 1.5) + elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: + self.learning_rate = min(1e-2, self.learning_rate * 1.5) + + # Update the learning rate for all GPUs + if self.is_multi_gpu: + lr_tensor = torch.tensor(self.learning_rate, device=self.device) + torch.distributed.broadcast(lr_tensor, src=0) + self.learning_rate = lr_tensor.item() + + # Update the learning rate for all parameter groups + for param_group in self.optimizer.param_groups: + param_group["lr"] = self.learning_rate + + # Surrogate loss + ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) + surrogate = -torch.squeeze(advantages_batch) * ratio + surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( + ratio, 1.0 - self.clip_param, 1.0 + self.clip_param + ) + surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() + + # Value function loss + if self.use_clipped_value_loss: + value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( + -self.clip_param, self.clip_param + ) + value_losses = (value_batch - returns_batch).pow(2) + value_losses_clipped = (value_clipped - returns_batch).pow(2) + value_loss = torch.max(value_losses, value_losses_clipped).mean() + else: + value_loss = (returns_batch - value_batch).pow(2).mean() + + loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() + + # Symmetry loss + if self.symmetry: + # obtain the symmetric actions + # if we did augmentation before then we don't need to augment again + if not self.symmetry["use_data_augmentation"]: + data_augmentation_func = self.symmetry["data_augmentation_func"] + obs_batch, _ = data_augmentation_func(obs=obs_batch, actions=None, env=self.symmetry["_env"]) + # compute number of augmentations per sample + num_aug = int(obs_batch.shape[0] / original_batch_size) + + # actions predicted by the actor for symmetrically-augmented observations + mean_actions_batch = self.policy.act_inference(obs_batch.detach().clone()) + + # compute the symmetrically augmented actions + # note: we are assuming the first augmentation is the original one. + # We do not use the action_batch from earlier since that action was sampled from the distribution. + # However, the symmetry loss is computed using the mean of the distribution. + action_mean_orig = mean_actions_batch[:original_batch_size] + _, actions_mean_symm_batch = data_augmentation_func( + obs=None, actions=action_mean_orig, env=self.symmetry["_env"] + ) + + # compute the loss (we skip the first augmentation as it is the original one) + mse_loss = torch.nn.MSELoss() + symmetry_loss = mse_loss( + mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] + ) + # add the loss to the total loss + if self.symmetry["use_mirror_loss"]: + loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss + else: + symmetry_loss = symmetry_loss.detach() + + # Random Network Distillation loss + # TODO: Move this processing to inside RND module. + if self.rnd: + # extract the rnd_state + # TODO: Check if we still need torch no grad. It is just an affine transformation. + with torch.no_grad(): + rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size]) + rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch) + # predict the embedding and the target + predicted_embedding = self.rnd.predictor(rnd_state_batch) + target_embedding = self.rnd.target(rnd_state_batch).detach() + # compute the loss as the mean squared error + mseloss = torch.nn.MSELoss() + rnd_loss = mseloss(predicted_embedding, target_embedding) + + expert_states = sample_amp_expert + policy_states = sample_amp_policy + + with torch.no_grad(): + expert_states = self.amp_normalizer.normalize_torch(expert_states.to(self.device), self.device) + policy_states = self.amp_normalizer.normalize_torch(policy_states, self.device) + + contact_phase_push = obs_batch['critic'][:, -4] + mask_push = contact_phase_push == 1. + + if mask_push.any(): + policy_d = self.discriminator(policy_states.flatten(1)) + expert_states = expert_states.to(self.device) + expert_d = self.discriminator(expert_states.flatten(1)) + + expert_loss = torch.nn.MSELoss()(expert_d, torch.ones(expert_d.size(), device=self.device)) + policy_loss = torch.nn.MSELoss()(policy_d, -1 * torch.ones(policy_d.size(), device=self.device)) + amp_loss = 0.5 * (expert_loss + policy_loss) + + # grad penalty + grad_pen_loss = self.discriminator.compute_grad_pen(expert_states, lambda_=5) + else: + amp_loss = torch.tensor(0.0, device=self.device) + grad_pen_loss = torch.tensor(0.0, device=self.device) + expert_loss = torch.tensor(0.0, device=self.device) + policy_loss = torch.tensor(0.0, device=self.device) + + loss += (amp_loss + grad_pen_loss) + self.amp_normalizer.update(policy_states.cpu().numpy()) + self.amp_normalizer.update(expert_states.cpu().numpy()) + + # Compute the gradients + # -- For PPO + self.optimizer.zero_grad() + loss.backward() + # -- For RND + if self.rnd: + self.rnd_optimizer.zero_grad() # type: ignore + rnd_loss.backward() + + # Collect gradients from all GPUs + if self.is_multi_gpu: + self.reduce_parameters() + + # Apply the gradients + # -- For PPO + nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) + self.optimizer.step() + # -- For RND + if self.rnd_optimizer: + self.rnd_optimizer.step() + + # Store the losses + mean_value_loss += value_loss.item() + mean_surrogate_loss += surrogate_loss.item() + mean_entropy += entropy_batch.mean().item() + mean_amp_loss += amp_loss.item() + mean_grad_pen_loss += grad_pen_loss.item() + mean_policy_pred += policy_loss.mean().item() + mean_expert_pred += expert_loss.mean().item() + # -- RND loss + if mean_rnd_loss is not None: + mean_rnd_loss += rnd_loss.item() + # -- Symmetry loss + if mean_symmetry_loss is not None: + mean_symmetry_loss += symmetry_loss.item() + + # -- For PPO + num_updates = self.num_learning_epochs * self.num_mini_batches + mean_value_loss /= num_updates + mean_surrogate_loss /= num_updates + mean_entropy /= num_updates + mean_amp_loss /= num_updates + mean_grad_pen_loss /= num_updates + mean_policy_pred /= num_updates + mean_expert_pred /= num_updates + + # -- For RND + if mean_rnd_loss is not None: + mean_rnd_loss /= num_updates + # -- For Symmetry + if mean_symmetry_loss is not None: + mean_symmetry_loss /= num_updates + # -- Clear the storage + self.storage.clear() + + # construct the loss dictionary + loss_dict = { + "value_function": mean_value_loss, + "surrogate": mean_surrogate_loss, + "entropy": mean_entropy, + "amp": mean_amp_loss, + "amp_grad_pen": mean_grad_pen_loss, + "amp_policy_pred": mean_policy_pred, + "amp_expert_pred": mean_expert_pred, + } + if self.rnd: + loss_dict["rnd"] = mean_rnd_loss + if self.symmetry: + loss_dict["symmetry"] = mean_symmetry_loss + + return loss_dict + + """ + Helper functions + """ + + def broadcast_parameters(self): + """Broadcast model parameters to all GPUs.""" + # obtain the model parameters on current GPU + model_params = [self.policy.state_dict()] + if self.rnd: + model_params.append(self.rnd.predictor.state_dict()) + # broadcast the model parameters + torch.distributed.broadcast_object_list(model_params, src=0) + # load the model parameters on all GPUs from source GPU + self.policy.load_state_dict(model_params[0]) + if self.rnd: + self.rnd.predictor.load_state_dict(model_params[1]) + + def reduce_parameters(self): + """Collect gradients from all GPUs and average them. + + This function is called after the backward pass to synchronize the gradients across all GPUs. + """ + # Create a tensor to store the gradients + grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] + if self.rnd: + grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None] + all_grads = torch.cat(grads) + + # Average the gradients across all GPUs + torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) + all_grads /= self.gpu_world_size + + # Get all parameters + all_params = self.policy.parameters() + if self.rnd: + all_params = chain(all_params, self.rnd.parameters()) + + # Update the gradients for all parameters with the reduced gradients + offset = 0 + for param in all_params: + if param.grad is not None: + numel = param.numel() + # copy data back from shared buffer + param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) + # update the offset for the next parameter + offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py new file mode 100644 index 0000000..3a86e00 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py @@ -0,0 +1,185 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch +import torch.nn as nn + +from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent +from rsl_rl.storage import RolloutStorage +from rsl_rl.utils import resolve_optimizer + + +class Distillation: + """Distillation algorithm for training a student model to mimic a teacher model.""" + + policy: StudentTeacher | StudentTeacherRecurrent + """The student teacher model.""" + + def __init__( + self, + policy, + num_learning_epochs=1, + gradient_length=15, + learning_rate=1e-3, + max_grad_norm=None, + loss_type="mse", + optimizer="adam", + device="cpu", + # Distributed training parameters + multi_gpu_cfg: dict | None = None, + ): + # device-related parameters + self.device = device + self.is_multi_gpu = multi_gpu_cfg is not None + # Multi-GPU parameters + if multi_gpu_cfg is not None: + self.gpu_global_rank = multi_gpu_cfg["global_rank"] + self.gpu_world_size = multi_gpu_cfg["world_size"] + else: + self.gpu_global_rank = 0 + self.gpu_world_size = 1 + + # distillation components + self.policy = policy + self.policy.to(self.device) + self.storage = None # initialized later + + # initialize the optimizer + self.optimizer = resolve_optimizer(optimizer)(self.policy.parameters(), lr=learning_rate) + + # initialize the transition + self.transition = RolloutStorage.Transition() + self.last_hidden_states = None + + # distillation parameters + self.num_learning_epochs = num_learning_epochs + self.gradient_length = gradient_length + self.learning_rate = learning_rate + self.max_grad_norm = max_grad_norm + + # initialize the loss function + loss_fn_dict = { + "mse": nn.functional.mse_loss, + "huber": nn.functional.huber_loss, + } + if loss_type in loss_fn_dict: + self.loss_fn = loss_fn_dict[loss_type] + else: + raise ValueError(f"Unknown loss type: {loss_type}. Supported types are: {list(loss_fn_dict.keys())}") + + self.num_updates = 0 + + def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): + # create rollout storage + self.storage = RolloutStorage( + training_type, + num_envs, + num_transitions_per_env, + obs, + actions_shape, + self.device, + ) + + def act(self, obs): + # compute the actions + self.transition.actions = self.policy.act(obs).detach() + self.transition.privileged_actions = self.policy.evaluate(obs).detach() + # record the observations + self.transition.observations = obs + return self.transition.actions + + def process_env_step(self, obs, rewards, dones, extras): + # update the normalizers + self.policy.update_normalization(obs) + + # record the rewards and dones + self.transition.rewards = rewards + self.transition.dones = dones + # record the transition + self.storage.add_transitions(self.transition) + self.transition.clear() + self.policy.reset(dones) + + def update(self): + self.num_updates += 1 + mean_behavior_loss = 0 + loss = 0 + cnt = 0 + + for epoch in range(self.num_learning_epochs): + self.policy.reset(hidden_states=self.last_hidden_states) + self.policy.detach_hidden_states() + for obs, _, privileged_actions, dones in self.storage.generator(): + + # inference the student for gradient computation + actions = self.policy.act_inference(obs) + + # behavior cloning loss + behavior_loss = self.loss_fn(actions, privileged_actions) + + # total loss + loss = loss + behavior_loss + mean_behavior_loss += behavior_loss.item() + cnt += 1 + + # gradient step + if cnt % self.gradient_length == 0: + self.optimizer.zero_grad() + loss.backward() + if self.is_multi_gpu: + self.reduce_parameters() + if self.max_grad_norm: + nn.utils.clip_grad_norm_(self.policy.student.parameters(), self.max_grad_norm) + self.optimizer.step() + self.policy.detach_hidden_states() + loss = 0 + + # reset dones + self.policy.reset(dones.view(-1)) + self.policy.detach_hidden_states(dones.view(-1)) + + mean_behavior_loss /= cnt + self.storage.clear() + self.last_hidden_states = self.policy.get_hidden_states() + self.policy.detach_hidden_states() + + # construct the loss dictionary + loss_dict = {"behavior": mean_behavior_loss} + + return loss_dict + + """ + Helper functions + """ + + def broadcast_parameters(self): + """Broadcast model parameters to all GPUs.""" + # obtain the model parameters on current GPU + model_params = [self.policy.state_dict()] + # broadcast the model parameters + torch.distributed.broadcast_object_list(model_params, src=0) + # load the model parameters on all GPUs from source GPU + self.policy.load_state_dict(model_params[0]) + + def reduce_parameters(self): + """Collect gradients from all GPUs and average them. + + This function is called after the backward pass to synchronize the gradients across all GPUs. + """ + # Create a tensor to store the gradients + grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] + all_grads = torch.cat(grads) + # Average the gradients across all GPUs + torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) + all_grads /= self.gpu_world_size + # Update the gradients for all parameters with the reduced gradients + offset = 0 + for param in self.policy.parameters(): + if param.grad is not None: + numel = param.numel() + # copy data back from shared buffer + param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) + # update the offset for the next parameter + offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py new file mode 100644 index 0000000..6c21fc5 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py @@ -0,0 +1,469 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.optim as optim +from itertools import chain + +from rsl_rl.modules import ActorCritic +from rsl_rl.modules.rnd import RandomNetworkDistillation +from rsl_rl.storage import RolloutStorage +from rsl_rl.utils import string_to_callable + + +class PPO: + """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" + + policy: ActorCritic + """The actor critic module.""" + + def __init__( + self, + policy, + num_learning_epochs=5, + num_mini_batches=4, + clip_param=0.2, + gamma=0.99, + lam=0.95, + value_loss_coef=1.0, + entropy_coef=0.01, + learning_rate=0.001, + max_grad_norm=1.0, + use_clipped_value_loss=True, + schedule="adaptive", + desired_kl=0.01, + device="cpu", + normalize_advantage_per_mini_batch=False, + # RND parameters + rnd_cfg: dict | None = None, + # Symmetry parameters + symmetry_cfg: dict | None = None, + # Distributed training parameters + multi_gpu_cfg: dict | None = None, + ): + # device-related parameters + self.device = device + self.is_multi_gpu = multi_gpu_cfg is not None + # Multi-GPU parameters + if multi_gpu_cfg is not None: + self.gpu_global_rank = multi_gpu_cfg["global_rank"] + self.gpu_world_size = multi_gpu_cfg["world_size"] + else: + self.gpu_global_rank = 0 + self.gpu_world_size = 1 + + # RND components + if rnd_cfg is not None: + # Extract parameters used in ppo + rnd_lr = rnd_cfg.pop("learning_rate", 1e-3) + # Create RND module + self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) + # Create RND optimizer + params = self.rnd.predictor.parameters() + self.rnd_optimizer = optim.Adam(params, lr=rnd_lr) + else: + self.rnd = None + self.rnd_optimizer = None + + # Symmetry components + if symmetry_cfg is not None: + # Check if symmetry is enabled + use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] + # Print that we are not using symmetry + if not use_symmetry: + print("Symmetry not used for learning. We will use it for logging instead.") + # If function is a string then resolve it to a function + if isinstance(symmetry_cfg["data_augmentation_func"], str): + symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) + # Check valid configuration + if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): + raise ValueError( + "Data augmentation enabled but the function is not callable:" + f" {symmetry_cfg['data_augmentation_func']}" + ) + # Store symmetry configuration + self.symmetry = symmetry_cfg + else: + self.symmetry = None + + # PPO components + self.policy = policy + self.policy.to(self.device) + # Create optimizer + self.optimizer = optim.Adam(self.policy.parameters(), lr=learning_rate) + # Create rollout storage + self.storage: RolloutStorage = None # type: ignore + self.transition = RolloutStorage.Transition() + + # PPO parameters + self.clip_param = clip_param + self.num_learning_epochs = num_learning_epochs + self.num_mini_batches = num_mini_batches + self.value_loss_coef = value_loss_coef + self.entropy_coef = entropy_coef + self.gamma = gamma + self.lam = lam + self.max_grad_norm = max_grad_norm + self.use_clipped_value_loss = use_clipped_value_loss + self.desired_kl = desired_kl + self.schedule = schedule + self.learning_rate = learning_rate + self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch + + def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): + # create rollout storage + self.storage = RolloutStorage( + training_type, + num_envs, + num_transitions_per_env, + obs, + actions_shape, + self.device, + ) + + def act(self, obs): + if self.policy.is_recurrent: + self.transition.hidden_states = self.policy.get_hidden_states() + # compute the actions and values + self.transition.actions = self.policy.act(obs).detach() + self.transition.values = self.policy.evaluate(obs).detach() + self.transition.actions_log_prob = self.policy.get_actions_log_prob(self.transition.actions).detach() + self.transition.action_mean = self.policy.action_mean.detach() + self.transition.action_sigma = self.policy.action_std.detach() + # need to record obs before env.step() + self.transition.observations = obs + return self.transition.actions + + def process_env_step(self, obs, rewards, dones, extras): + # update the normalizers + self.policy.update_normalization(obs) + if self.rnd: + self.rnd.update_normalization(obs) + + # Record the rewards and dones + # Note: we clone here because later on we bootstrap the rewards based on timeouts + self.transition.rewards = rewards.clone() + self.transition.dones = dones + + # Compute the intrinsic rewards and add to extrinsic rewards + if self.rnd: + # Compute the intrinsic rewards + self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs) + # Add intrinsic rewards to extrinsic rewards + self.transition.rewards += self.intrinsic_rewards + + # Bootstrapping on time outs + if "time_outs" in extras: + self.transition.rewards += self.gamma * torch.squeeze( + self.transition.values * extras["time_outs"].unsqueeze(1).to(self.device), 1 + ) + + # record the transition + self.storage.add_transitions(self.transition) + self.transition.clear() + self.policy.reset(dones) + + def compute_returns(self, obs): + # compute value for the last step + last_values = self.policy.evaluate(obs).detach() + self.storage.compute_returns( + last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch + ) + + def update(self): # noqa: C901 + mean_value_loss = 0 + mean_surrogate_loss = 0 + mean_entropy = 0 + # -- RND loss + if self.rnd: + mean_rnd_loss = 0 + else: + mean_rnd_loss = None + # -- Symmetry loss + if self.symmetry: + mean_symmetry_loss = 0 + else: + mean_symmetry_loss = None + + # generator for mini batches + if self.policy.is_recurrent: + generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + else: + generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + + # iterate over batches + for ( + obs_batch, + actions_batch, + target_values_batch, + advantages_batch, + returns_batch, + old_actions_log_prob_batch, + old_mu_batch, + old_sigma_batch, + hid_states_batch, + masks_batch, + ) in generator: + + # number of augmentations per sample + # we start with 1 and increase it if we use symmetry augmentation + num_aug = 1 + # original batch size + # we assume policy group is always there and needs augmentation + original_batch_size = obs_batch.batch_size[0] + + # check if we should normalize advantages per mini batch + if self.normalize_advantage_per_mini_batch: + with torch.no_grad(): + advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8) + + # Perform symmetric augmentation + if self.symmetry and self.symmetry["use_data_augmentation"]: + # augmentation using symmetry + data_augmentation_func = self.symmetry["data_augmentation_func"] + # returned shape: [batch_size * num_aug, ...] + obs_batch, actions_batch = data_augmentation_func( + obs=obs_batch, + actions=actions_batch, + env=self.symmetry["_env"], + ) + # compute number of augmentations per sample + # we assume policy group is always there and needs augmentation + num_aug = int(obs_batch.batch_size[0] / original_batch_size) + # repeat the rest of the batch + # -- actor + old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) + # -- critic + target_values_batch = target_values_batch.repeat(num_aug, 1) + advantages_batch = advantages_batch.repeat(num_aug, 1) + returns_batch = returns_batch.repeat(num_aug, 1) + + # Recompute actions log prob and entropy for current batch of transitions + # Note: we need to do this because we updated the policy with the new parameters + # -- actor + self.policy.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) + actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch) + # -- critic + value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) + # -- entropy + # we only keep the entropy of the first augmentation (the original one) + mu_batch = self.policy.action_mean[:original_batch_size] + sigma_batch = self.policy.action_std[:original_batch_size] + entropy_batch = self.policy.entropy[:original_batch_size] + + # KL + if self.desired_kl is not None and self.schedule == "adaptive": + with torch.inference_mode(): + kl = torch.sum( + torch.log(sigma_batch / old_sigma_batch + 1.0e-5) + + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) + / (2.0 * torch.square(sigma_batch)) + - 0.5, + axis=-1, + ) + kl_mean = torch.mean(kl) + + # Reduce the KL divergence across all GPUs + if self.is_multi_gpu: + torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM) + kl_mean /= self.gpu_world_size + + # Update the learning rate + # Perform this adaptation only on the main process + # TODO: Is this needed? If KL-divergence is the "same" across all GPUs, + # then the learning rate should be the same across all GPUs. + if self.gpu_global_rank == 0: + if kl_mean > self.desired_kl * 2.0: + self.learning_rate = max(1e-5, self.learning_rate / 1.5) + elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: + self.learning_rate = min(1e-2, self.learning_rate * 1.5) + + # Update the learning rate for all GPUs + if self.is_multi_gpu: + lr_tensor = torch.tensor(self.learning_rate, device=self.device) + torch.distributed.broadcast(lr_tensor, src=0) + self.learning_rate = lr_tensor.item() + + # Update the learning rate for all parameter groups + for param_group in self.optimizer.param_groups: + param_group["lr"] = self.learning_rate + + # Surrogate loss + ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) + surrogate = -torch.squeeze(advantages_batch) * ratio + surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( + ratio, 1.0 - self.clip_param, 1.0 + self.clip_param + ) + surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() + + # Value function loss + if self.use_clipped_value_loss: + value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( + -self.clip_param, self.clip_param + ) + value_losses = (value_batch - returns_batch).pow(2) + value_losses_clipped = (value_clipped - returns_batch).pow(2) + value_loss = torch.max(value_losses, value_losses_clipped).mean() + else: + value_loss = (returns_batch - value_batch).pow(2).mean() + + loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() + + # Symmetry loss + if self.symmetry: + # obtain the symmetric actions + # if we did augmentation before then we don't need to augment again + if not self.symmetry["use_data_augmentation"]: + data_augmentation_func = self.symmetry["data_augmentation_func"] + obs_batch, _ = data_augmentation_func(obs=obs_batch, actions=None, env=self.symmetry["_env"]) + # compute number of augmentations per sample + num_aug = int(obs_batch.shape[0] / original_batch_size) + + # actions predicted by the actor for symmetrically-augmented observations + mean_actions_batch = self.policy.act_inference(obs_batch.detach().clone()) + + # compute the symmetrically augmented actions + # note: we are assuming the first augmentation is the original one. + # We do not use the action_batch from earlier since that action was sampled from the distribution. + # However, the symmetry loss is computed using the mean of the distribution. + action_mean_orig = mean_actions_batch[:original_batch_size] + _, actions_mean_symm_batch = data_augmentation_func( + obs=None, actions=action_mean_orig, env=self.symmetry["_env"] + ) + + # compute the loss (we skip the first augmentation as it is the original one) + mse_loss = torch.nn.MSELoss() + symmetry_loss = mse_loss( + mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] + ) + # add the loss to the total loss + if self.symmetry["use_mirror_loss"]: + loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss + else: + symmetry_loss = symmetry_loss.detach() + + # Random Network Distillation loss + # TODO: Move this processing to inside RND module. + if self.rnd: + # extract the rnd_state + # TODO: Check if we still need torch no grad. It is just an affine transformation. + with torch.no_grad(): + rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size]) + rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch) + # predict the embedding and the target + predicted_embedding = self.rnd.predictor(rnd_state_batch) + target_embedding = self.rnd.target(rnd_state_batch).detach() + # compute the loss as the mean squared error + mseloss = torch.nn.MSELoss() + rnd_loss = mseloss(predicted_embedding, target_embedding) + + # Compute the gradients + # -- For PPO + self.optimizer.zero_grad() + loss.backward() + # -- For RND + if self.rnd: + self.rnd_optimizer.zero_grad() # type: ignore + rnd_loss.backward() + + # Collect gradients from all GPUs + if self.is_multi_gpu: + self.reduce_parameters() + + # Apply the gradients + # -- For PPO + nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) + self.optimizer.step() + # -- For RND + if self.rnd_optimizer: + self.rnd_optimizer.step() + + # Store the losses + mean_value_loss += value_loss.item() + mean_surrogate_loss += surrogate_loss.item() + mean_entropy += entropy_batch.mean().item() + # -- RND loss + if mean_rnd_loss is not None: + mean_rnd_loss += rnd_loss.item() + # -- Symmetry loss + if mean_symmetry_loss is not None: + mean_symmetry_loss += symmetry_loss.item() + + # -- For PPO + num_updates = self.num_learning_epochs * self.num_mini_batches + mean_value_loss /= num_updates + mean_surrogate_loss /= num_updates + mean_entropy /= num_updates + # -- For RND + if mean_rnd_loss is not None: + mean_rnd_loss /= num_updates + # -- For Symmetry + if mean_symmetry_loss is not None: + mean_symmetry_loss /= num_updates + # -- Clear the storage + self.storage.clear() + + # construct the loss dictionary + loss_dict = { + "value_function": mean_value_loss, + "surrogate": mean_surrogate_loss, + "entropy": mean_entropy, + } + if self.rnd: + loss_dict["rnd"] = mean_rnd_loss + if self.symmetry: + loss_dict["symmetry"] = mean_symmetry_loss + + return loss_dict + + """ + Helper functions + """ + + def broadcast_parameters(self): + """Broadcast model parameters to all GPUs.""" + # obtain the model parameters on current GPU + model_params = [self.policy.state_dict()] + if self.rnd: + model_params.append(self.rnd.predictor.state_dict()) + # broadcast the model parameters + torch.distributed.broadcast_object_list(model_params, src=0) + # load the model parameters on all GPUs from source GPU + self.policy.load_state_dict(model_params[0]) + if self.rnd: + self.rnd.predictor.load_state_dict(model_params[1]) + + def reduce_parameters(self): + """Collect gradients from all GPUs and average them. + + This function is called after the backward pass to synchronize the gradients across all GPUs. + """ + # Create a tensor to store the gradients + grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] + if self.rnd: + grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None] + all_grads = torch.cat(grads) + + # Average the gradients across all GPUs + torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) + all_grads /= self.gpu_world_size + + # Get all parameters + all_params = self.policy.parameters() + if self.rnd: + all_params = chain(all_params, self.rnd.parameters()) + + # Update the gradients for all parameters with the reduced gradients + offset = 0 + for param in all_params: + if param.grad is not None: + numel = param.numel() + # copy data back from shared buffer + param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) + # update the offset for the next parameter + offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/__init__.py new file mode 100644 index 0000000..ebc2e20 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Main module for the rsl_rl package.""" diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py new file mode 100644 index 0000000..0ef9163 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Implementation of different RL agents.""" + +from .distillation import Distillation +from .ppo import PPO +from .amp_ppo import AMP_PPO +__all__ = ["PPO", "Distillation", "AMP_PPO"] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py new file mode 100644 index 0000000..a2164eb --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py @@ -0,0 +1,571 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations +from torch._tensor import Tensor +from torch._tensor import Tensor +from typing import Any +from copy import deepcopy + +import torch +import torch.nn as nn +import torch.optim as optim +from itertools import chain + +from rsl_rl.modules import ActorCritic +from rsl_rl.modules.rnd import RandomNetworkDistillation +from rsl_rl.storage import RolloutStorage, ReplayBufferMulti +from rsl_rl.utils import string_to_callable + + +class AMP_PPO: + """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" + + policy: ActorCritic + """The actor critic module.""" + + def __init__( + self, + policy, + discriminator, + amp_data, + amp_normalizer, + amp_num_frames=1, + amp_replay_buffer_size=100000, + num_learning_epochs=5, + num_mini_batches=4, + clip_param=0.2, + gamma=0.99, + lam=0.95, + value_loss_coef=1.0, + entropy_coef=0.01, + learning_rate=0.001, + max_grad_norm=1.0, + use_clipped_value_loss=True, + schedule="adaptive", + desired_kl=0.01, + device="cpu", + normalize_advantage_per_mini_batch=False, + # RND parameters + rnd_cfg: dict | None = None, + # Symmetry parameters + symmetry_cfg: dict | None = None, + # Distributed training parameters + multi_gpu_cfg: dict | None = None, + ): + # device-related parameters + self.device = device + self.is_multi_gpu = multi_gpu_cfg is not None + # Multi-GPU parameters + if multi_gpu_cfg is not None: + self.gpu_global_rank = multi_gpu_cfg["global_rank"] + self.gpu_world_size = multi_gpu_cfg["world_size"] + else: + self.gpu_global_rank = 0 + self.gpu_world_size = 1 + + # RND components + if rnd_cfg is not None: + # Extract parameters used in ppo + rnd_lr = rnd_cfg.pop("learning_rate", 1e-3) + # Create RND module + self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) + # Create RND optimizer + params = self.rnd.predictor.parameters() + self.rnd_optimizer = optim.Adam(params, lr=rnd_lr) + else: + self.rnd = None + self.rnd_optimizer = None + + # Symmetry components + if symmetry_cfg is not None: + # Check if symmetry is enabled + use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] + # Print that we are not using symmetry + if not use_symmetry: + print("Symmetry not used for learning. We will use it for logging instead.") + # If function is a string then resolve it to a function + if isinstance(symmetry_cfg["data_augmentation_func"], str): + symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) + # Check valid configuration + if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): + raise ValueError( + "Data augmentation enabled but the function is not callable:" + f" {symmetry_cfg['data_augmentation_func']}" + ) + # Store symmetry configuration + self.symmetry = symmetry_cfg + else: + self.symmetry = None + + ## AMP components + self.discriminator = discriminator + self.discriminator.to(self.device) + + self.amp_storage = ReplayBufferMulti(discriminator.state_dim, amp_replay_buffer_size, amp_num_frames, device) + self.amp_data = amp_data + self.amp_normalizer = amp_normalizer + + # PPO components + self.policy = policy + self.policy.to(self.device) + + # Create rollout storage + self.storage: RolloutStorage = None # type: ignore + self.transition = RolloutStorage.Transition() + self.amp_transition = RolloutStorage.Transition() + params = [ + {'params': self.policy.parameters(), 'name': 'policy'}, + ] + + params.append({ + 'params': self.discriminator.trunk.parameters(), + 'weight_decay': 10e-4, + 'name': f'amp_trunk' + }) + params.append({ + 'params': self.discriminator.amp_linear.parameters(), + 'weight_decay': 10e-2, + 'name': f'amp_head' + }) + + # Create optimizer + self.optimizer = optim.Adam(params, lr=learning_rate) + + # PPO parameters + self.clip_param = clip_param + self.num_learning_epochs = num_learning_epochs + self.num_mini_batches = num_mini_batches + self.value_loss_coef = value_loss_coef + self.entropy_coef = entropy_coef + self.gamma = gamma + self.lam = lam + self.max_grad_norm = max_grad_norm + self.use_clipped_value_loss = use_clipped_value_loss + self.desired_kl = desired_kl + self.schedule = schedule + self.learning_rate = learning_rate + self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch + + + def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): + # create rollout storage + self.storage = RolloutStorage( + training_type, + num_envs, + num_transitions_per_env, + obs, + actions_shape, + self.device, + ) + + def act(self, obs, amp_obs): + if self.policy.is_recurrent: + self.transition.hidden_states = self.policy.get_hidden_states() + # compute the actions and values + self.transition.actions = self.policy.act(obs).detach() + self.transition.values = self.policy.evaluate(obs).detach() + self.transition.actions_log_prob = self.policy.get_actions_log_prob(self.transition.actions).detach() + self.transition.action_mean = self.policy.action_mean.detach() + self.transition.action_sigma = self.policy.action_std.detach() + # need to record obs before env.step() + self.transition.observations = obs + self.amp_transition.observations = amp_obs + return self.transition.actions + + def process_env_step(self, obs, rewards, dones, extras,amp_obs, amp_obs_frames=None): + # update the normalizers + self.policy.update_normalization(obs) + if self.rnd: + self.rnd.update_normalization(obs) + + # Record the rewards and dones + # Note: we clone here because later on we bootstrap the rewards based on timeouts + self.transition.rewards = rewards.clone() + self.transition.dones = dones + + # Compute the intrinsic rewards and add to extrinsic rewards + if self.rnd: + # Compute the intrinsic rewards + self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs) + # Add intrinsic rewards to extrinsic rewards + self.transition.rewards += self.intrinsic_rewards + + # Bootstrapping on time outs + if "time_outs" in extras: + self.transition.rewards += self.gamma * torch.squeeze( + self.transition.values * extras["time_outs"].unsqueeze(1).to(self.device), 1 + ) + + if amp_obs_frames is not None: + self.amp_storage.insert(amp_obs_frames) + else: + self.amp_storage.insert(self.amp_transition.observations, amp_obs) + + # record the transition + self.storage.add_transitions(self.transition) + self.transition.clear() + self.amp_transition.clear() + self.policy.reset(dones) + + def compute_returns(self, obs): + # compute value for the last step + last_values = self.policy.evaluate(obs).detach() + self.storage.compute_returns( + last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch + ) + + def update(self): # noqa: C901 + mean_value_loss = 0 + mean_surrogate_loss = 0 + mean_entropy = 0 + mean_amp_loss = 0 + mean_grad_pen_loss = 0 + mean_policy_pred = 0 + mean_expert_pred = 0 + # -- RND loss + if self.rnd: + mean_rnd_loss = 0 + else: + mean_rnd_loss = None + # -- Symmetry loss + if self.symmetry: + mean_symmetry_loss = 0 + else: + mean_symmetry_loss = None + + # generator for mini batches + if self.policy.is_recurrent: + generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + else: + generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + + + amp_policy_generator = self.amp_storage.feed_forward_generator( + self.num_learning_epochs * self.num_mini_batches, + self.storage.num_envs * self.storage.num_transitions_per_env // self.num_mini_batches, + ) + + amp_expert_generator = self.amp_data.feed_forward_generator_23dof_multi( + self.num_learning_epochs * self.num_mini_batches, + self.storage.num_envs * self.storage.num_transitions_per_env // self.num_mini_batches, + ) + + # iterate over batches + for sample, sample_amp_policy, sample_amp_expert in zip(generator, amp_policy_generator, amp_expert_generator): + ( + obs_batch, + actions_batch, + target_values_batch, + advantages_batch, + returns_batch, + old_actions_log_prob_batch, + old_mu_batch, + old_sigma_batch, + hid_states_batch, + masks_batch, + ) = sample + + # number of augmentations per sample + # we start with 1 and increase it if we use symmetry augmentation + num_aug = 1 + # original batch size + # we assume policy group is always there and needs augmentation + original_batch_size = obs_batch.batch_size[0] + + # check if we should normalize advantages per mini batch + if self.normalize_advantage_per_mini_batch: + with torch.no_grad(): + advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8) + + # Perform symmetric augmentation + if self.symmetry and self.symmetry["use_data_augmentation"]: + # augmentation using symmetry + data_augmentation_func = self.symmetry["data_augmentation_func"] + # returned shape: [batch_size * num_aug, ...] + obs_batch, actions_batch = data_augmentation_func( + obs=obs_batch, + actions=actions_batch, + env=self.symmetry["_env"], + ) + # compute number of augmentations per sample + # we assume policy group is always there and needs augmentation + num_aug = int(obs_batch.batch_size[0] / original_batch_size) + # repeat the rest of the batch + # -- actor + old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) + # -- critic + target_values_batch = target_values_batch.repeat(num_aug, 1) + advantages_batch = advantages_batch.repeat(num_aug, 1) + returns_batch = returns_batch.repeat(num_aug, 1) + + # Recompute actions log prob and entropy for current batch of transitions + # Note: we need to do this because we updated the policy with the new parameters + # -- actor + self.policy.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) + actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch) + # -- critic + value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) + # -- entropy + # we only keep the entropy of the first augmentation (the original one) + mu_batch = self.policy.action_mean[:original_batch_size] + sigma_batch = self.policy.action_std[:original_batch_size] + entropy_batch = self.policy.entropy[:original_batch_size] + + # KL + if self.desired_kl is not None and self.schedule == "adaptive": + with torch.inference_mode(): + kl = torch.sum( + torch.log(sigma_batch / old_sigma_batch + 1.0e-5) + + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) + / (2.0 * torch.square(sigma_batch)) + - 0.5, + axis=-1, + ) + kl_mean = torch.mean(kl) + + # Reduce the KL divergence across all GPUs + if self.is_multi_gpu: + torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM) + kl_mean /= self.gpu_world_size + + # Update the learning rate + # Perform this adaptation only on the main process + # TODO: Is this needed? If KL-divergence is the "same" across all GPUs, + # then the learning rate should be the same across all GPUs. + if self.gpu_global_rank == 0: + if kl_mean > self.desired_kl * 2.0: + self.learning_rate = max(1e-5, self.learning_rate / 1.5) + elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: + self.learning_rate = min(1e-2, self.learning_rate * 1.5) + + # Update the learning rate for all GPUs + if self.is_multi_gpu: + lr_tensor = torch.tensor(self.learning_rate, device=self.device) + torch.distributed.broadcast(lr_tensor, src=0) + self.learning_rate = lr_tensor.item() + + # Update the learning rate for all parameter groups + for param_group in self.optimizer.param_groups: + param_group["lr"] = self.learning_rate + + # Surrogate loss + ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) + surrogate = -torch.squeeze(advantages_batch) * ratio + surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( + ratio, 1.0 - self.clip_param, 1.0 + self.clip_param + ) + surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() + + # Value function loss + if self.use_clipped_value_loss: + value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( + -self.clip_param, self.clip_param + ) + value_losses = (value_batch - returns_batch).pow(2) + value_losses_clipped = (value_clipped - returns_batch).pow(2) + value_loss = torch.max(value_losses, value_losses_clipped).mean() + else: + value_loss = (returns_batch - value_batch).pow(2).mean() + + loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() + + # Symmetry loss + if self.symmetry: + # obtain the symmetric actions + # if we did augmentation before then we don't need to augment again + if not self.symmetry["use_data_augmentation"]: + data_augmentation_func = self.symmetry["data_augmentation_func"] + obs_batch, _ = data_augmentation_func(obs=obs_batch, actions=None, env=self.symmetry["_env"]) + # compute number of augmentations per sample + num_aug = int(obs_batch.shape[0] / original_batch_size) + + # actions predicted by the actor for symmetrically-augmented observations + mean_actions_batch = self.policy.act_inference(obs_batch.detach().clone()) + + # compute the symmetrically augmented actions + # note: we are assuming the first augmentation is the original one. + # We do not use the action_batch from earlier since that action was sampled from the distribution. + # However, the symmetry loss is computed using the mean of the distribution. + action_mean_orig = mean_actions_batch[:original_batch_size] + _, actions_mean_symm_batch = data_augmentation_func( + obs=None, actions=action_mean_orig, env=self.symmetry["_env"] + ) + + # compute the loss (we skip the first augmentation as it is the original one) + mse_loss = torch.nn.MSELoss() + symmetry_loss = mse_loss( + mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] + ) + # add the loss to the total loss + if self.symmetry["use_mirror_loss"]: + loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss + else: + symmetry_loss = symmetry_loss.detach() + + # Random Network Distillation loss + # TODO: Move this processing to inside RND module. + if self.rnd: + # extract the rnd_state + # TODO: Check if we still need torch no grad. It is just an affine transformation. + with torch.no_grad(): + rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size]) + rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch) + # predict the embedding and the target + predicted_embedding = self.rnd.predictor(rnd_state_batch) + target_embedding = self.rnd.target(rnd_state_batch).detach() + # compute the loss as the mean squared error + mseloss = torch.nn.MSELoss() + rnd_loss = mseloss(predicted_embedding, target_embedding) + + expert_states = sample_amp_expert + policy_states = sample_amp_policy + + with torch.no_grad(): + expert_states = self.amp_normalizer.normalize_torch(expert_states.to(self.device), self.device) + policy_states = self.amp_normalizer.normalize_torch(policy_states, self.device) + + contact_phase_push = obs_batch['critic'][:, -4] + mask_push = contact_phase_push == 1. + + if mask_push.any(): + policy_d = self.discriminator(policy_states.flatten(1)) + expert_states = expert_states.to(self.device) + expert_d = self.discriminator(expert_states.flatten(1)) + + expert_loss = torch.nn.MSELoss()(expert_d, torch.ones(expert_d.size(), device=self.device)) + policy_loss = torch.nn.MSELoss()(policy_d, -1 * torch.ones(policy_d.size(), device=self.device)) + amp_loss = 0.5 * (expert_loss + policy_loss) + + # grad penalty + grad_pen_loss = self.discriminator.compute_grad_pen(expert_states, lambda_=5) + else: + amp_loss = torch.tensor(0.0, device=self.device) + grad_pen_loss = torch.tensor(0.0, device=self.device) + expert_loss = torch.tensor(0.0, device=self.device) + policy_loss = torch.tensor(0.0, device=self.device) + + loss += (amp_loss + grad_pen_loss) + self.amp_normalizer.update(policy_states.cpu().numpy()) + self.amp_normalizer.update(expert_states.cpu().numpy()) + + # Compute the gradients + # -- For PPO + self.optimizer.zero_grad() + loss.backward() + # -- For RND + if self.rnd: + self.rnd_optimizer.zero_grad() # type: ignore + rnd_loss.backward() + + # Collect gradients from all GPUs + if self.is_multi_gpu: + self.reduce_parameters() + + # Apply the gradients + # -- For PPO + nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) + self.optimizer.step() + # -- For RND + if self.rnd_optimizer: + self.rnd_optimizer.step() + + # Store the losses + mean_value_loss += value_loss.item() + mean_surrogate_loss += surrogate_loss.item() + mean_entropy += entropy_batch.mean().item() + mean_amp_loss += amp_loss.item() + mean_grad_pen_loss += grad_pen_loss.item() + mean_policy_pred += policy_loss.mean().item() + mean_expert_pred += expert_loss.mean().item() + # -- RND loss + if mean_rnd_loss is not None: + mean_rnd_loss += rnd_loss.item() + # -- Symmetry loss + if mean_symmetry_loss is not None: + mean_symmetry_loss += symmetry_loss.item() + + # -- For PPO + num_updates = self.num_learning_epochs * self.num_mini_batches + mean_value_loss /= num_updates + mean_surrogate_loss /= num_updates + mean_entropy /= num_updates + mean_amp_loss /= num_updates + mean_grad_pen_loss /= num_updates + mean_policy_pred /= num_updates + mean_expert_pred /= num_updates + + # -- For RND + if mean_rnd_loss is not None: + mean_rnd_loss /= num_updates + # -- For Symmetry + if mean_symmetry_loss is not None: + mean_symmetry_loss /= num_updates + # -- Clear the storage + self.storage.clear() + + # construct the loss dictionary + loss_dict = { + "value_function": mean_value_loss, + "surrogate": mean_surrogate_loss, + "entropy": mean_entropy, + "amp": mean_amp_loss, + "amp_grad_pen": mean_grad_pen_loss, + "amp_policy_pred": mean_policy_pred, + "amp_expert_pred": mean_expert_pred, + } + if self.rnd: + loss_dict["rnd"] = mean_rnd_loss + if self.symmetry: + loss_dict["symmetry"] = mean_symmetry_loss + + return loss_dict + + """ + Helper functions + """ + + def broadcast_parameters(self): + """Broadcast model parameters to all GPUs.""" + # obtain the model parameters on current GPU + model_params = [self.policy.state_dict()] + if self.rnd: + model_params.append(self.rnd.predictor.state_dict()) + # broadcast the model parameters + torch.distributed.broadcast_object_list(model_params, src=0) + # load the model parameters on all GPUs from source GPU + self.policy.load_state_dict(model_params[0]) + if self.rnd: + self.rnd.predictor.load_state_dict(model_params[1]) + + def reduce_parameters(self): + """Collect gradients from all GPUs and average them. + + This function is called after the backward pass to synchronize the gradients across all GPUs. + """ + # Create a tensor to store the gradients + grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] + if self.rnd: + grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None] + all_grads = torch.cat(grads) + + # Average the gradients across all GPUs + torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) + all_grads /= self.gpu_world_size + + # Get all parameters + all_params = self.policy.parameters() + if self.rnd: + all_params = chain(all_params, self.rnd.parameters()) + + # Update the gradients for all parameters with the reduced gradients + offset = 0 + for param in all_params: + if param.grad is not None: + numel = param.numel() + # copy data back from shared buffer + param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) + # update the offset for the next parameter + offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py new file mode 100644 index 0000000..3a86e00 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py @@ -0,0 +1,185 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import torch +import torch.nn as nn + +from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent +from rsl_rl.storage import RolloutStorage +from rsl_rl.utils import resolve_optimizer + + +class Distillation: + """Distillation algorithm for training a student model to mimic a teacher model.""" + + policy: StudentTeacher | StudentTeacherRecurrent + """The student teacher model.""" + + def __init__( + self, + policy, + num_learning_epochs=1, + gradient_length=15, + learning_rate=1e-3, + max_grad_norm=None, + loss_type="mse", + optimizer="adam", + device="cpu", + # Distributed training parameters + multi_gpu_cfg: dict | None = None, + ): + # device-related parameters + self.device = device + self.is_multi_gpu = multi_gpu_cfg is not None + # Multi-GPU parameters + if multi_gpu_cfg is not None: + self.gpu_global_rank = multi_gpu_cfg["global_rank"] + self.gpu_world_size = multi_gpu_cfg["world_size"] + else: + self.gpu_global_rank = 0 + self.gpu_world_size = 1 + + # distillation components + self.policy = policy + self.policy.to(self.device) + self.storage = None # initialized later + + # initialize the optimizer + self.optimizer = resolve_optimizer(optimizer)(self.policy.parameters(), lr=learning_rate) + + # initialize the transition + self.transition = RolloutStorage.Transition() + self.last_hidden_states = None + + # distillation parameters + self.num_learning_epochs = num_learning_epochs + self.gradient_length = gradient_length + self.learning_rate = learning_rate + self.max_grad_norm = max_grad_norm + + # initialize the loss function + loss_fn_dict = { + "mse": nn.functional.mse_loss, + "huber": nn.functional.huber_loss, + } + if loss_type in loss_fn_dict: + self.loss_fn = loss_fn_dict[loss_type] + else: + raise ValueError(f"Unknown loss type: {loss_type}. Supported types are: {list(loss_fn_dict.keys())}") + + self.num_updates = 0 + + def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): + # create rollout storage + self.storage = RolloutStorage( + training_type, + num_envs, + num_transitions_per_env, + obs, + actions_shape, + self.device, + ) + + def act(self, obs): + # compute the actions + self.transition.actions = self.policy.act(obs).detach() + self.transition.privileged_actions = self.policy.evaluate(obs).detach() + # record the observations + self.transition.observations = obs + return self.transition.actions + + def process_env_step(self, obs, rewards, dones, extras): + # update the normalizers + self.policy.update_normalization(obs) + + # record the rewards and dones + self.transition.rewards = rewards + self.transition.dones = dones + # record the transition + self.storage.add_transitions(self.transition) + self.transition.clear() + self.policy.reset(dones) + + def update(self): + self.num_updates += 1 + mean_behavior_loss = 0 + loss = 0 + cnt = 0 + + for epoch in range(self.num_learning_epochs): + self.policy.reset(hidden_states=self.last_hidden_states) + self.policy.detach_hidden_states() + for obs, _, privileged_actions, dones in self.storage.generator(): + + # inference the student for gradient computation + actions = self.policy.act_inference(obs) + + # behavior cloning loss + behavior_loss = self.loss_fn(actions, privileged_actions) + + # total loss + loss = loss + behavior_loss + mean_behavior_loss += behavior_loss.item() + cnt += 1 + + # gradient step + if cnt % self.gradient_length == 0: + self.optimizer.zero_grad() + loss.backward() + if self.is_multi_gpu: + self.reduce_parameters() + if self.max_grad_norm: + nn.utils.clip_grad_norm_(self.policy.student.parameters(), self.max_grad_norm) + self.optimizer.step() + self.policy.detach_hidden_states() + loss = 0 + + # reset dones + self.policy.reset(dones.view(-1)) + self.policy.detach_hidden_states(dones.view(-1)) + + mean_behavior_loss /= cnt + self.storage.clear() + self.last_hidden_states = self.policy.get_hidden_states() + self.policy.detach_hidden_states() + + # construct the loss dictionary + loss_dict = {"behavior": mean_behavior_loss} + + return loss_dict + + """ + Helper functions + """ + + def broadcast_parameters(self): + """Broadcast model parameters to all GPUs.""" + # obtain the model parameters on current GPU + model_params = [self.policy.state_dict()] + # broadcast the model parameters + torch.distributed.broadcast_object_list(model_params, src=0) + # load the model parameters on all GPUs from source GPU + self.policy.load_state_dict(model_params[0]) + + def reduce_parameters(self): + """Collect gradients from all GPUs and average them. + + This function is called after the backward pass to synchronize the gradients across all GPUs. + """ + # Create a tensor to store the gradients + grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] + all_grads = torch.cat(grads) + # Average the gradients across all GPUs + torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) + all_grads /= self.gpu_world_size + # Update the gradients for all parameters with the reduced gradients + offset = 0 + for param in self.policy.parameters(): + if param.grad is not None: + numel = param.numel() + # copy data back from shared buffer + param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) + # update the offset for the next parameter + offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py new file mode 100644 index 0000000..6c21fc5 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py @@ -0,0 +1,469 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.optim as optim +from itertools import chain + +from rsl_rl.modules import ActorCritic +from rsl_rl.modules.rnd import RandomNetworkDistillation +from rsl_rl.storage import RolloutStorage +from rsl_rl.utils import string_to_callable + + +class PPO: + """Proximal Policy Optimization algorithm (https://arxiv.org/abs/1707.06347).""" + + policy: ActorCritic + """The actor critic module.""" + + def __init__( + self, + policy, + num_learning_epochs=5, + num_mini_batches=4, + clip_param=0.2, + gamma=0.99, + lam=0.95, + value_loss_coef=1.0, + entropy_coef=0.01, + learning_rate=0.001, + max_grad_norm=1.0, + use_clipped_value_loss=True, + schedule="adaptive", + desired_kl=0.01, + device="cpu", + normalize_advantage_per_mini_batch=False, + # RND parameters + rnd_cfg: dict | None = None, + # Symmetry parameters + symmetry_cfg: dict | None = None, + # Distributed training parameters + multi_gpu_cfg: dict | None = None, + ): + # device-related parameters + self.device = device + self.is_multi_gpu = multi_gpu_cfg is not None + # Multi-GPU parameters + if multi_gpu_cfg is not None: + self.gpu_global_rank = multi_gpu_cfg["global_rank"] + self.gpu_world_size = multi_gpu_cfg["world_size"] + else: + self.gpu_global_rank = 0 + self.gpu_world_size = 1 + + # RND components + if rnd_cfg is not None: + # Extract parameters used in ppo + rnd_lr = rnd_cfg.pop("learning_rate", 1e-3) + # Create RND module + self.rnd = RandomNetworkDistillation(device=self.device, **rnd_cfg) + # Create RND optimizer + params = self.rnd.predictor.parameters() + self.rnd_optimizer = optim.Adam(params, lr=rnd_lr) + else: + self.rnd = None + self.rnd_optimizer = None + + # Symmetry components + if symmetry_cfg is not None: + # Check if symmetry is enabled + use_symmetry = symmetry_cfg["use_data_augmentation"] or symmetry_cfg["use_mirror_loss"] + # Print that we are not using symmetry + if not use_symmetry: + print("Symmetry not used for learning. We will use it for logging instead.") + # If function is a string then resolve it to a function + if isinstance(symmetry_cfg["data_augmentation_func"], str): + symmetry_cfg["data_augmentation_func"] = string_to_callable(symmetry_cfg["data_augmentation_func"]) + # Check valid configuration + if symmetry_cfg["use_data_augmentation"] and not callable(symmetry_cfg["data_augmentation_func"]): + raise ValueError( + "Data augmentation enabled but the function is not callable:" + f" {symmetry_cfg['data_augmentation_func']}" + ) + # Store symmetry configuration + self.symmetry = symmetry_cfg + else: + self.symmetry = None + + # PPO components + self.policy = policy + self.policy.to(self.device) + # Create optimizer + self.optimizer = optim.Adam(self.policy.parameters(), lr=learning_rate) + # Create rollout storage + self.storage: RolloutStorage = None # type: ignore + self.transition = RolloutStorage.Transition() + + # PPO parameters + self.clip_param = clip_param + self.num_learning_epochs = num_learning_epochs + self.num_mini_batches = num_mini_batches + self.value_loss_coef = value_loss_coef + self.entropy_coef = entropy_coef + self.gamma = gamma + self.lam = lam + self.max_grad_norm = max_grad_norm + self.use_clipped_value_loss = use_clipped_value_loss + self.desired_kl = desired_kl + self.schedule = schedule + self.learning_rate = learning_rate + self.normalize_advantage_per_mini_batch = normalize_advantage_per_mini_batch + + def init_storage(self, training_type, num_envs, num_transitions_per_env, obs, actions_shape): + # create rollout storage + self.storage = RolloutStorage( + training_type, + num_envs, + num_transitions_per_env, + obs, + actions_shape, + self.device, + ) + + def act(self, obs): + if self.policy.is_recurrent: + self.transition.hidden_states = self.policy.get_hidden_states() + # compute the actions and values + self.transition.actions = self.policy.act(obs).detach() + self.transition.values = self.policy.evaluate(obs).detach() + self.transition.actions_log_prob = self.policy.get_actions_log_prob(self.transition.actions).detach() + self.transition.action_mean = self.policy.action_mean.detach() + self.transition.action_sigma = self.policy.action_std.detach() + # need to record obs before env.step() + self.transition.observations = obs + return self.transition.actions + + def process_env_step(self, obs, rewards, dones, extras): + # update the normalizers + self.policy.update_normalization(obs) + if self.rnd: + self.rnd.update_normalization(obs) + + # Record the rewards and dones + # Note: we clone here because later on we bootstrap the rewards based on timeouts + self.transition.rewards = rewards.clone() + self.transition.dones = dones + + # Compute the intrinsic rewards and add to extrinsic rewards + if self.rnd: + # Compute the intrinsic rewards + self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs) + # Add intrinsic rewards to extrinsic rewards + self.transition.rewards += self.intrinsic_rewards + + # Bootstrapping on time outs + if "time_outs" in extras: + self.transition.rewards += self.gamma * torch.squeeze( + self.transition.values * extras["time_outs"].unsqueeze(1).to(self.device), 1 + ) + + # record the transition + self.storage.add_transitions(self.transition) + self.transition.clear() + self.policy.reset(dones) + + def compute_returns(self, obs): + # compute value for the last step + last_values = self.policy.evaluate(obs).detach() + self.storage.compute_returns( + last_values, self.gamma, self.lam, normalize_advantage=not self.normalize_advantage_per_mini_batch + ) + + def update(self): # noqa: C901 + mean_value_loss = 0 + mean_surrogate_loss = 0 + mean_entropy = 0 + # -- RND loss + if self.rnd: + mean_rnd_loss = 0 + else: + mean_rnd_loss = None + # -- Symmetry loss + if self.symmetry: + mean_symmetry_loss = 0 + else: + mean_symmetry_loss = None + + # generator for mini batches + if self.policy.is_recurrent: + generator = self.storage.recurrent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + else: + generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs) + + # iterate over batches + for ( + obs_batch, + actions_batch, + target_values_batch, + advantages_batch, + returns_batch, + old_actions_log_prob_batch, + old_mu_batch, + old_sigma_batch, + hid_states_batch, + masks_batch, + ) in generator: + + # number of augmentations per sample + # we start with 1 and increase it if we use symmetry augmentation + num_aug = 1 + # original batch size + # we assume policy group is always there and needs augmentation + original_batch_size = obs_batch.batch_size[0] + + # check if we should normalize advantages per mini batch + if self.normalize_advantage_per_mini_batch: + with torch.no_grad(): + advantages_batch = (advantages_batch - advantages_batch.mean()) / (advantages_batch.std() + 1e-8) + + # Perform symmetric augmentation + if self.symmetry and self.symmetry["use_data_augmentation"]: + # augmentation using symmetry + data_augmentation_func = self.symmetry["data_augmentation_func"] + # returned shape: [batch_size * num_aug, ...] + obs_batch, actions_batch = data_augmentation_func( + obs=obs_batch, + actions=actions_batch, + env=self.symmetry["_env"], + ) + # compute number of augmentations per sample + # we assume policy group is always there and needs augmentation + num_aug = int(obs_batch.batch_size[0] / original_batch_size) + # repeat the rest of the batch + # -- actor + old_actions_log_prob_batch = old_actions_log_prob_batch.repeat(num_aug, 1) + # -- critic + target_values_batch = target_values_batch.repeat(num_aug, 1) + advantages_batch = advantages_batch.repeat(num_aug, 1) + returns_batch = returns_batch.repeat(num_aug, 1) + + # Recompute actions log prob and entropy for current batch of transitions + # Note: we need to do this because we updated the policy with the new parameters + # -- actor + self.policy.act(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[0]) + actions_log_prob_batch = self.policy.get_actions_log_prob(actions_batch) + # -- critic + value_batch = self.policy.evaluate(obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1]) + # -- entropy + # we only keep the entropy of the first augmentation (the original one) + mu_batch = self.policy.action_mean[:original_batch_size] + sigma_batch = self.policy.action_std[:original_batch_size] + entropy_batch = self.policy.entropy[:original_batch_size] + + # KL + if self.desired_kl is not None and self.schedule == "adaptive": + with torch.inference_mode(): + kl = torch.sum( + torch.log(sigma_batch / old_sigma_batch + 1.0e-5) + + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) + / (2.0 * torch.square(sigma_batch)) + - 0.5, + axis=-1, + ) + kl_mean = torch.mean(kl) + + # Reduce the KL divergence across all GPUs + if self.is_multi_gpu: + torch.distributed.all_reduce(kl_mean, op=torch.distributed.ReduceOp.SUM) + kl_mean /= self.gpu_world_size + + # Update the learning rate + # Perform this adaptation only on the main process + # TODO: Is this needed? If KL-divergence is the "same" across all GPUs, + # then the learning rate should be the same across all GPUs. + if self.gpu_global_rank == 0: + if kl_mean > self.desired_kl * 2.0: + self.learning_rate = max(1e-5, self.learning_rate / 1.5) + elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0: + self.learning_rate = min(1e-2, self.learning_rate * 1.5) + + # Update the learning rate for all GPUs + if self.is_multi_gpu: + lr_tensor = torch.tensor(self.learning_rate, device=self.device) + torch.distributed.broadcast(lr_tensor, src=0) + self.learning_rate = lr_tensor.item() + + # Update the learning rate for all parameter groups + for param_group in self.optimizer.param_groups: + param_group["lr"] = self.learning_rate + + # Surrogate loss + ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch)) + surrogate = -torch.squeeze(advantages_batch) * ratio + surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp( + ratio, 1.0 - self.clip_param, 1.0 + self.clip_param + ) + surrogate_loss = torch.max(surrogate, surrogate_clipped).mean() + + # Value function loss + if self.use_clipped_value_loss: + value_clipped = target_values_batch + (value_batch - target_values_batch).clamp( + -self.clip_param, self.clip_param + ) + value_losses = (value_batch - returns_batch).pow(2) + value_losses_clipped = (value_clipped - returns_batch).pow(2) + value_loss = torch.max(value_losses, value_losses_clipped).mean() + else: + value_loss = (returns_batch - value_batch).pow(2).mean() + + loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() + + # Symmetry loss + if self.symmetry: + # obtain the symmetric actions + # if we did augmentation before then we don't need to augment again + if not self.symmetry["use_data_augmentation"]: + data_augmentation_func = self.symmetry["data_augmentation_func"] + obs_batch, _ = data_augmentation_func(obs=obs_batch, actions=None, env=self.symmetry["_env"]) + # compute number of augmentations per sample + num_aug = int(obs_batch.shape[0] / original_batch_size) + + # actions predicted by the actor for symmetrically-augmented observations + mean_actions_batch = self.policy.act_inference(obs_batch.detach().clone()) + + # compute the symmetrically augmented actions + # note: we are assuming the first augmentation is the original one. + # We do not use the action_batch from earlier since that action was sampled from the distribution. + # However, the symmetry loss is computed using the mean of the distribution. + action_mean_orig = mean_actions_batch[:original_batch_size] + _, actions_mean_symm_batch = data_augmentation_func( + obs=None, actions=action_mean_orig, env=self.symmetry["_env"] + ) + + # compute the loss (we skip the first augmentation as it is the original one) + mse_loss = torch.nn.MSELoss() + symmetry_loss = mse_loss( + mean_actions_batch[original_batch_size:], actions_mean_symm_batch.detach()[original_batch_size:] + ) + # add the loss to the total loss + if self.symmetry["use_mirror_loss"]: + loss += self.symmetry["mirror_loss_coeff"] * symmetry_loss + else: + symmetry_loss = symmetry_loss.detach() + + # Random Network Distillation loss + # TODO: Move this processing to inside RND module. + if self.rnd: + # extract the rnd_state + # TODO: Check if we still need torch no grad. It is just an affine transformation. + with torch.no_grad(): + rnd_state_batch = self.rnd.get_rnd_state(obs_batch[:original_batch_size]) + rnd_state_batch = self.rnd.state_normalizer(rnd_state_batch) + # predict the embedding and the target + predicted_embedding = self.rnd.predictor(rnd_state_batch) + target_embedding = self.rnd.target(rnd_state_batch).detach() + # compute the loss as the mean squared error + mseloss = torch.nn.MSELoss() + rnd_loss = mseloss(predicted_embedding, target_embedding) + + # Compute the gradients + # -- For PPO + self.optimizer.zero_grad() + loss.backward() + # -- For RND + if self.rnd: + self.rnd_optimizer.zero_grad() # type: ignore + rnd_loss.backward() + + # Collect gradients from all GPUs + if self.is_multi_gpu: + self.reduce_parameters() + + # Apply the gradients + # -- For PPO + nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) + self.optimizer.step() + # -- For RND + if self.rnd_optimizer: + self.rnd_optimizer.step() + + # Store the losses + mean_value_loss += value_loss.item() + mean_surrogate_loss += surrogate_loss.item() + mean_entropy += entropy_batch.mean().item() + # -- RND loss + if mean_rnd_loss is not None: + mean_rnd_loss += rnd_loss.item() + # -- Symmetry loss + if mean_symmetry_loss is not None: + mean_symmetry_loss += symmetry_loss.item() + + # -- For PPO + num_updates = self.num_learning_epochs * self.num_mini_batches + mean_value_loss /= num_updates + mean_surrogate_loss /= num_updates + mean_entropy /= num_updates + # -- For RND + if mean_rnd_loss is not None: + mean_rnd_loss /= num_updates + # -- For Symmetry + if mean_symmetry_loss is not None: + mean_symmetry_loss /= num_updates + # -- Clear the storage + self.storage.clear() + + # construct the loss dictionary + loss_dict = { + "value_function": mean_value_loss, + "surrogate": mean_surrogate_loss, + "entropy": mean_entropy, + } + if self.rnd: + loss_dict["rnd"] = mean_rnd_loss + if self.symmetry: + loss_dict["symmetry"] = mean_symmetry_loss + + return loss_dict + + """ + Helper functions + """ + + def broadcast_parameters(self): + """Broadcast model parameters to all GPUs.""" + # obtain the model parameters on current GPU + model_params = [self.policy.state_dict()] + if self.rnd: + model_params.append(self.rnd.predictor.state_dict()) + # broadcast the model parameters + torch.distributed.broadcast_object_list(model_params, src=0) + # load the model parameters on all GPUs from source GPU + self.policy.load_state_dict(model_params[0]) + if self.rnd: + self.rnd.predictor.load_state_dict(model_params[1]) + + def reduce_parameters(self): + """Collect gradients from all GPUs and average them. + + This function is called after the backward pass to synchronize the gradients across all GPUs. + """ + # Create a tensor to store the gradients + grads = [param.grad.view(-1) for param in self.policy.parameters() if param.grad is not None] + if self.rnd: + grads += [param.grad.view(-1) for param in self.rnd.parameters() if param.grad is not None] + all_grads = torch.cat(grads) + + # Average the gradients across all GPUs + torch.distributed.all_reduce(all_grads, op=torch.distributed.ReduceOp.SUM) + all_grads /= self.gpu_world_size + + # Get all parameters + all_params = self.policy.parameters() + if self.rnd: + all_params = chain(all_params, self.rnd.parameters()) + + # Update the gradients for all parameters with the reduced gradients + offset = 0 + for param in all_params: + if param.grad is not None: + numel = param.numel() + # copy data back from shared buffer + param.grad.data.copy_(all_grads[offset : offset + numel].view_as(param.grad.data)) + # update the offset for the next parameter + offset += numel diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/__init__.py new file mode 100644 index 0000000..ab7c056 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Submodule defining the environment definitions.""" + +from .vec_env import VecEnv + +__all__ = ["VecEnv"] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/vec_env.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/vec_env.py new file mode 100644 index 0000000..ae471cf --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/vec_env.py @@ -0,0 +1,113 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from abc import ABC, abstractmethod +from tensordict import TensorDict + + +class VecEnv(ABC): + """Abstract class for a vectorized environment. + + The vectorized environment is a collection of environments that are synchronized. This means that + the same type of action is applied to all environments and the same type of observation is returned from all + environments. + """ + + num_envs: int + """Number of environments.""" + + num_actions: int + """Number of actions.""" + + max_episode_length: int | torch.Tensor + + max_episode_length_s: float + """Maximum episode length. + + The maximum episode length can be a scalar or a tensor. If it is a scalar, it is the same for all environments. + If it is a tensor, it is the maximum episode length for each environment. This is useful for dynamic episode + lengths. + """ + + episode_length_buf: torch.Tensor + """Buffer for current episode lengths.""" + + device: torch.device | str + """Device to use.""" + + cfg: dict | object + """Configuration object.""" + + reset_env_ids: torch.Tensor | None = None + + contact_phase: torch.Tensor | None = None + """ + Operations. + """ + + @abstractmethod + def get_observations(self) -> TensorDict: + """Return the current observations. + + Returns: + observations (TensorDict): Observations from the environment. + """ + raise NotImplementedError + + @abstractmethod + def get_amp_observations(self) -> TensorDict: + """Return the current AMP observations. + + Returns: + observations (TensorDict): Observations from the environment. + """ + raise NotImplementedError + + @abstractmethod + def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]: + """Apply input action to the environment. + + Args: + actions (torch.Tensor): Input actions to apply. Shape: (num_envs, num_actions) + + Returns: + observations (TensorDict): Observations from the environment. + rewards (torch.Tensor): Rewards from the environment. Shape: (num_envs,) + dones (torch.Tensor): Done flags from the environment. Shape: (num_envs,) + extras (dict): Extra information from the environment. + + Observations: + + The observations TensorDict usually contains multiple observation groups. The `obs_groups` + dictionary of the runner configuration specifies which observation groups are used for which + purpose, i.e., it maps the available observation groups to observation sets. The observation sets + (keys of the `obs_groups` dictionary) currently used by rsl_rl are: + + - "policy": Specified observation groups are used as input to the actor/student network. + - "critic": Specified observation groups are used as input to the critic network. + - "teacher": Specified observation groups are used as input to the teacher network. + - "rnd_state": Specified observation groups are used as input to the RND network. + + Incomplete or incorrect configurations are handled in the `resolve_obs_groups()` function in + `rsl_rl/utils/utils.py`. + + Extras: + + The extras dictionary includes metrics such as the episode reward, episode length, etc. The following + dictionary keys are used by rsl_rl: + + - "time_outs" (torch.Tensor): Timeouts for the environments. These correspond to terminations that + happen due to time limits and not due to the environment reaching a terminal state. This is useful + for environments that have a fixed episode length. + + - "log" (dict[str, float | torch.Tensor]): Additional information for logging and debugging purposes. + The key should be a string and start with "/" for namespacing. The value can be a scalar or a + tensor. If it is a tensor, the mean of the tensor is used for logging. + """ + raise NotImplementedError + \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/__init__.py new file mode 100644 index 0000000..9afe1a2 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Definitions for neural-network components for RL-agents.""" + +from .actor_critic import ActorCritic +from .actor_critic_recurrent import ActorCriticRecurrent +from .rnd import * +from .student_teacher import StudentTeacher +from .student_teacher_recurrent import StudentTeacherRecurrent +from .symmetry import * +from .discriminator_multi import DiscriminatorMulti +__all__ = [ + "ActorCritic", + "ActorCriticRecurrent", + "StudentTeacher", + "StudentTeacherRecurrent", + "DiscriminatorMulti", +] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py new file mode 100644 index 0000000..0efc36b --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py @@ -0,0 +1,195 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +from torch.distributions import Normal + +from rsl_rl.networks import MLP, EmpiricalNormalization + + +class ActorCritic(nn.Module): + is_recurrent = False + + def __init__( + self, + obs, + obs_groups, + num_actions, + actor_obs_normalization=False, + critic_obs_normalization=False, + actor_hidden_dims=[256, 256, 256], + critic_hidden_dims=[256, 256, 256], + activation="elu", + init_noise_std=1.0, + noise_std_type: str = "scalar", + state_dependent_std=False, + **kwargs, + ): + if kwargs: + print( + "ActorCritic.__init__ got unexpected arguments, which will be ignored: " + + str([key for key in kwargs.keys()]) + ) + super().__init__() + + # get the observation dimensions + self.obs_groups = obs_groups + num_actor_obs = 0 + for obs_group in obs_groups["policy"]: + assert len(obs[obs_group].shape) == 2, "The ActorCritic module only supports 1D observations." + num_actor_obs += obs[obs_group].shape[-1] + num_critic_obs = 0 + for obs_group in obs_groups["critic"]: + assert len(obs[obs_group].shape) == 2, "The ActorCritic module only supports 1D observations." + num_critic_obs += obs[obs_group].shape[-1] + + self.state_dependent_std = state_dependent_std + # actor + if self.state_dependent_std: + self.actor = MLP(num_actor_obs, [2, num_actions], actor_hidden_dims, activation) + else: + self.actor = MLP(num_actor_obs, num_actions, actor_hidden_dims, activation) + # actor observation normalization + self.actor_obs_normalization = actor_obs_normalization + if actor_obs_normalization: + self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs) + else: + self.actor_obs_normalizer = torch.nn.Identity() + print(f"Actor MLP: {self.actor}") + + # critic + self.critic = MLP(num_critic_obs, 1, critic_hidden_dims, activation) + # critic observation normalization + self.critic_obs_normalization = critic_obs_normalization + if critic_obs_normalization: + self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs) + else: + self.critic_obs_normalizer = torch.nn.Identity() + print(f"Critic MLP: {self.critic}") + + # Action noise + self.noise_std_type = noise_std_type + if self.state_dependent_std: + torch.nn.init.zeros_(self.actor[-2].weight[num_actions:]) + if self.noise_std_type == "scalar": + torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std) + elif self.noise_std_type == "log": + torch.nn.init.constant_( + self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7)) + ) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + else: + if self.noise_std_type == "scalar": + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + elif self.noise_std_type == "log": + self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + + # Action distribution (populated in update_distribution) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args(False) + + def reset(self, dones=None): + pass + + def forward(self): + raise NotImplementedError + + @property + def action_mean(self): + return self.distribution.mean + + @property + def action_std(self): + return self.distribution.stddev + + @property + def entropy(self): + return self.distribution.entropy().sum(dim=-1) + + def update_distribution(self, obs): + if self.state_dependent_std: + # compute mean and standard deviation + mean_and_std = self.actor(obs) + if self.noise_std_type == "scalar": + mean, std = torch.unbind(mean_and_std, dim=-2) + elif self.noise_std_type == "log": + mean, log_std = torch.unbind(mean_and_std, dim=-2) + std = torch.exp(log_std) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + else: + # compute mean + mean = self.actor(obs) + # compute standard deviation + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + # create distribution + self.distribution = Normal(mean, std) + + def act(self, obs, **kwargs): + obs = self.get_actor_obs(obs) + obs = self.actor_obs_normalizer(obs) + self.update_distribution(obs) + return self.distribution.sample() + + def act_inference(self, obs): + obs = self.get_actor_obs(obs) + obs = self.actor_obs_normalizer(obs) + return self.actor(obs) + + def evaluate(self, obs, **kwargs): + obs = self.get_critic_obs(obs) + obs = self.critic_obs_normalizer(obs) + return self.critic(obs) + + def get_actor_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["policy"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_critic_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["critic"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_actions_log_prob(self, actions): + return self.distribution.log_prob(actions).sum(dim=-1) + + def update_normalization(self, obs): + if self.actor_obs_normalization: + actor_obs = self.get_actor_obs(obs) + self.actor_obs_normalizer.update(actor_obs) + if self.critic_obs_normalization: + critic_obs = self.get_critic_obs(obs) + self.critic_obs_normalizer.update(critic_obs) + + def load_state_dict(self, state_dict, strict=True): + """Load the parameters of the actor-critic model. + + Args: + state_dict (dict): State dictionary of the model. + strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this + module's state_dict() function. + + Returns: + bool: Whether this training resumes a previous training. This flag is used by the `load()` function of + `OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation). + """ + + super().load_state_dict(state_dict, strict=strict) + return True # training resumes diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py new file mode 100644 index 0000000..bba46ca --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py @@ -0,0 +1,218 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +import warnings +from torch.distributions import Normal + +from rsl_rl.networks import MLP, EmpiricalNormalization, Memory + + +class ActorCriticRecurrent(nn.Module): + is_recurrent = True + + def __init__( + self, + obs, + obs_groups, + num_actions, + actor_obs_normalization=False, + critic_obs_normalization=False, + actor_hidden_dims=[256, 256, 256], + critic_hidden_dims=[256, 256, 256], + activation="elu", + init_noise_std=1.0, + noise_std_type: str = "scalar", + state_dependent_std=False, + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=1, + **kwargs, + ): + if "rnn_hidden_size" in kwargs: + warnings.warn( + "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. " + "Please use `rnn_hidden_dim` instead.", + DeprecationWarning, + ) + if rnn_hidden_dim == 256: # Only override if the new argument is at its default + rnn_hidden_dim = kwargs.pop("rnn_hidden_size") + if kwargs: + print( + "ActorCriticRecurrent.__init__ got unexpected arguments, which will be ignored: " + str(kwargs.keys()), + ) + super().__init__() + + # get the observation dimensions + self.obs_groups = obs_groups + num_actor_obs = 0 + for obs_group in obs_groups["policy"]: + assert len(obs[obs_group].shape) == 2, "The ActorCriticRecurrent module only supports 1D observations." + num_actor_obs += obs[obs_group].shape[-1] + num_critic_obs = 0 + for obs_group in obs_groups["critic"]: + assert len(obs[obs_group].shape) == 2, "The ActorCriticRecurrent module only supports 1D observations." + num_critic_obs += obs[obs_group].shape[-1] + + self.state_dependent_std = state_dependent_std + # actor + self.memory_a = Memory(num_actor_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) + if self.state_dependent_std: + self.actor = MLP(rnn_hidden_dim, [2, num_actions], actor_hidden_dims, activation) + else: + self.actor = MLP(rnn_hidden_dim, num_actions, actor_hidden_dims, activation) + + # actor observation normalization + self.actor_obs_normalization = actor_obs_normalization + if actor_obs_normalization: + self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs) + else: + self.actor_obs_normalizer = torch.nn.Identity() + print(f"Actor RNN: {self.memory_a}") + print(f"Actor MLP: {self.actor}") + + # critic + self.memory_c = Memory(num_critic_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) + self.critic = MLP(rnn_hidden_dim, 1, critic_hidden_dims, activation) + # critic observation normalization + self.critic_obs_normalization = critic_obs_normalization + if critic_obs_normalization: + self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs) + else: + self.critic_obs_normalizer = torch.nn.Identity() + print(f"Critic RNN: {self.memory_c}") + print(f"Critic MLP: {self.critic}") + + # Action noise + self.noise_std_type = noise_std_type + if self.state_dependent_std: + torch.nn.init.zeros_(self.actor[-2].weight[num_actions:]) + if self.noise_std_type == "scalar": + torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std) + elif self.noise_std_type == "log": + torch.nn.init.constant_( + self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7)) + ) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + else: + if self.noise_std_type == "scalar": + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + elif self.noise_std_type == "log": + self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + + # Action distribution (populated in update_distribution) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args(False) + + @property + def action_mean(self): + return self.distribution.mean + + @property + def action_std(self): + return self.distribution.stddev + + @property + def entropy(self): + return self.distribution.entropy().sum(dim=-1) + + def reset(self, dones=None): + self.memory_a.reset(dones) + self.memory_c.reset(dones) + + def forward(self): + raise NotImplementedError + + def update_distribution(self, obs): + if self.state_dependent_std: + # compute mean and standard deviation + mean_and_std = self.actor(obs) + if self.noise_std_type == "scalar": + mean, std = torch.unbind(mean_and_std, dim=-2) + elif self.noise_std_type == "log": + mean, log_std = torch.unbind(mean_and_std, dim=-2) + std = torch.exp(log_std) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + else: + # compute mean + mean = self.actor(obs) + # compute standard deviation + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + # create distribution + self.distribution = Normal(mean, std) + + def act(self, obs, masks=None, hidden_states=None): + obs = self.get_actor_obs(obs) + obs = self.actor_obs_normalizer(obs) + out_mem = self.memory_a(obs, masks, hidden_states).squeeze(0) + self.update_distribution(out_mem) + return self.distribution.sample() + + def act_inference(self, obs): + obs = self.get_actor_obs(obs) + obs = self.actor_obs_normalizer(obs) + out_mem = self.memory_a(obs).squeeze(0) + return self.actor(out_mem) + + def evaluate(self, obs, masks=None, hidden_states=None): + obs = self.get_critic_obs(obs) + obs = self.critic_obs_normalizer(obs) + out_mem = self.memory_c(obs, masks, hidden_states).squeeze(0) + return self.critic(out_mem) + + def get_actor_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["policy"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_critic_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["critic"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_actions_log_prob(self, actions): + return self.distribution.log_prob(actions).sum(dim=-1) + + def get_hidden_states(self): + return self.memory_a.hidden_states, self.memory_c.hidden_states + + def update_normalization(self, obs): + if self.actor_obs_normalization: + actor_obs = self.get_actor_obs(obs) + self.actor_obs_normalizer.update(actor_obs) + if self.critic_obs_normalization: + critic_obs = self.get_critic_obs(obs) + self.critic_obs_normalizer.update(critic_obs) + + def load_state_dict(self, state_dict, strict=True): + """Load the parameters of the actor-critic model. + + Args: + state_dict (dict): State dictionary of the model. + strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this + module's state_dict() function. + + Returns: + bool: Whether this training resumes a previous training. This flag is used by the `load()` function of + `OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation). + """ + + super().load_state_dict(state_dict, strict=strict) + return True diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py new file mode 100644 index 0000000..7899d23 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py @@ -0,0 +1,102 @@ +import torch +import torch.nn as nn +from torch import autograd +import torch.nn.utils.spectral_norm as spectral_norm + + +class DiscriminatorMulti(nn.Module): + def __init__( + self, state_dim, amp_reward_coef, hidden_layer_sizes, device, + num_frames=2, task_reward_lerp=0.0, use_lerp=True): + super(DiscriminatorMulti, self).__init__() + + self.device = device + self.state_dim = state_dim + self.use_lerp = use_lerp + self.num_frames = num_frames # 存储帧数参数 + + self.amp_reward_coef = amp_reward_coef + amp_layers = [] + + curr_in_dim = state_dim * num_frames + for hidden_dim in hidden_layer_sizes: + amp_layers.append(spectral_norm(nn.Linear(curr_in_dim, hidden_dim))) + amp_layers.append(nn.ReLU()) + curr_in_dim = hidden_dim + self.trunk = nn.Sequential(*amp_layers).to(device) + self.amp_linear = spectral_norm(nn.Linear(hidden_layer_sizes[-1], 1)).to(device) + + self.trunk.train() + self.amp_linear.train() + + self.task_reward_lerp = task_reward_lerp + + def forward(self, x): + h = self.trunk(x) + d = self.amp_linear(h) + return d + + def compute_grad_pen(self, + expert_states, # 改为接收多帧状态列表 + lambda_=10): + # 将多帧状态沿最后一个维度拼接 + expert_data = expert_states.flatten(1) + expert_data.requires_grad = True + + disc = self.amp_linear(self.trunk(expert_data)) + ones = torch.ones(disc.size(), device=disc.device) + grad = autograd.grad( + outputs=disc, inputs=expert_data, + grad_outputs=ones, create_graph=True, + retain_graph=True, only_inputs=True)[0] + + # Enforce that the grad norm approaches 0. + grad_pen = lambda_ * (grad.norm(2, dim=1) - 0).pow(2).mean() + return grad_pen + + + def get_disc_weights(self): + weights = [] + for m in self.trunk.modules(): + if isinstance(m, nn.Linear): + weights.append(torch.flatten(m.weight)) + + weights.append(torch.flatten(self.amp_linear.weight)) + return weights + + def get_disc_logit_weights(self): + return torch.flatten(self.amp_linear.weight) + + def predict_amp_reward( + self, states, # 改为接收多帧状态列表 + task_reward, normalizer=None): + """ + states: torch.Tensor, shape=(num_envs, num_frames, state_dim) + task_reward: torch.Tensor, shape=(num_envs, 1) + """ + # import ipdb; ipdb.set_trace() + with torch.no_grad(): + self.eval() + if normalizer is not None: + # 对每一帧状态进行归一化 + states = normalizer.normalize_torch(states, self.device) + + # 拼接多帧状态 + state_cat = states.flatten(1) + d = self.amp_linear(self.trunk(state_cat)) + disc_reward = self.amp_reward_coef * torch.clamp(1 - (1/4) * torch.square(d - 1), min=0) + + if self.use_lerp: + if self.task_reward_lerp > 0: + reward = self._lerp_reward(disc_reward, task_reward.unsqueeze(-1)) + self.train() + return reward.squeeze(), d, disc_reward.squeeze() * (1.0 - self.task_reward_lerp) + else: + disc_reward *= 0.02 + reward = task_reward.unsqueeze(-1) + disc_reward + self.train() + return reward.squeeze(), d, disc_reward.squeeze() + + def _lerp_reward(self, disc_r, task_r): + r = (1.0 - self.task_reward_lerp) * disc_r + self.task_reward_lerp * task_r + return r \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/rnd.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/rnd.py new file mode 100644 index 0000000..8e65c43 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/rnd.py @@ -0,0 +1,209 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn + +from rsl_rl.networks import MLP, EmpiricalDiscountedVariationNormalization, EmpiricalNormalization + + +class RandomNetworkDistillation(nn.Module): + """Implementation of Random Network Distillation (RND) [1] + + References: + .. [1] Burda, Yuri, et al. "Exploration by random network distillation." arXiv preprint arXiv:1810.12894 (2018). + """ + + def __init__( + self, + num_states: int, + obs_groups: dict, + num_outputs: int, + predictor_hidden_dims: list[int], + target_hidden_dims: list[int], + activation: str = "elu", + weight: float = 0.0, + state_normalization: bool = False, + reward_normalization: bool = False, + device: str = "cpu", + weight_schedule: dict | None = None, + ): + """Initialize the RND module. + + - If :attr:`state_normalization` is True, then the input state is normalized using an Empirical Normalization layer. + - If :attr:`reward_normalization` is True, then the intrinsic reward is normalized using an Empirical Discounted + Variation Normalization layer. + + .. note:: + If the hidden dimensions are -1 in the predictor and target networks configuration, then the number of states + is used as the hidden dimension. + + Args: + num_states: Number of states/inputs to the predictor and target networks. + num_outputs: Number of outputs (embedding size) of the predictor and target networks. + predictor_hidden_dims: List of hidden dimensions of the predictor network. + target_hidden_dims: List of hidden dimensions of the target network. + activation: Activation function. Defaults to "elu". + weight: Scaling factor of the intrinsic reward. Defaults to 0.0. + state_normalization: Whether to normalize the input state. Defaults to False. + reward_normalization: Whether to normalize the intrinsic reward. Defaults to False. + device: Device to use. Defaults to "cpu". + weight_schedule: The type of schedule to use for the RND weight parameter. + Defaults to None, in which case the weight parameter is constant. + It is a dictionary with the following keys: + + - "mode": The type of schedule to use for the RND weight parameter. + - "constant": Constant weight schedule. + - "step": Step weight schedule. + - "linear": Linear weight schedule. + + For the "step" weight schedule, the following parameters are required: + + - "final_step": The step at which the weight parameter is set to the final value. + - "final_value": The final value of the weight parameter. + + For the "linear" weight schedule, the following parameters are required: + - "initial_step": The step at which the weight parameter is set to the initial value. + - "final_step": The step at which the weight parameter is set to the final value. + - "final_value": The final value of the weight parameter. + """ + # initialize parent class + super().__init__() + + # Store parameters + self.num_states = num_states + self.obs_groups = obs_groups + self.num_outputs = num_outputs + self.initial_weight = weight + self.device = device + self.state_normalization = state_normalization + self.reward_normalization = reward_normalization + + # Normalization of input gates + if state_normalization: + self.state_normalizer = EmpiricalNormalization(shape=[self.num_states], until=1.0e8).to(self.device) + else: + self.state_normalizer = torch.nn.Identity() + # Normalization of intrinsic reward + if reward_normalization: + self.reward_normalizer = EmpiricalDiscountedVariationNormalization(shape=[], until=1.0e8).to(self.device) + else: + self.reward_normalizer = torch.nn.Identity() + + # counter for the number of updates + self.update_counter = 0 + + # resolve weight schedule + if weight_schedule is not None: + self.weight_scheduler_params = weight_schedule + self.weight_scheduler = getattr(self, f"_{weight_schedule['mode']}_weight_schedule") + else: + self.weight_scheduler = None + # Create network architecture + self.predictor = MLP(num_states, num_outputs, predictor_hidden_dims, activation).to(self.device) + self.target = MLP(num_states, num_outputs, target_hidden_dims, activation).to(self.device) + + # make target network not trainable + self.target.eval() + + def get_intrinsic_reward(self, obs) -> torch.Tensor: + # Note: the counter is updated number of env steps per learning iteration + self.update_counter += 1 + # Extract the rnd state from the observation + rnd_state = self.get_rnd_state(obs) + rnd_state = self.state_normalizer(rnd_state) + # Obtain the embedding of the rnd state from the target and predictor networks + target_embedding = self.target(rnd_state).detach() + predictor_embedding = self.predictor(rnd_state).detach() + # Compute the intrinsic reward as the distance between the embeddings + intrinsic_reward = torch.linalg.norm(target_embedding - predictor_embedding, dim=1) + # Normalize intrinsic reward + intrinsic_reward = self.reward_normalizer(intrinsic_reward) + + # Check the weight schedule + if self.weight_scheduler is not None: + self.weight = self.weight_scheduler(step=self.update_counter, **self.weight_scheduler_params) + else: + self.weight = self.initial_weight + # Scale intrinsic reward + intrinsic_reward *= self.weight + + return intrinsic_reward + + def forward(self, *args, **kwargs): + raise RuntimeError("Forward method is not implemented. Use get_intrinsic_reward instead.") + + def train(self, mode: bool = True): + # sets module into training mode + self.predictor.train(mode) + if self.state_normalization: + self.state_normalizer.train(mode) + if self.reward_normalization: + self.reward_normalizer.train(mode) + return self + + def eval(self): + return self.train(False) + + def get_rnd_state(self, obs): + obs_list = [] + for obs_group in self.obs_groups["rnd_state"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def update_normalization(self, obs): + # Normalize the state + if self.state_normalization: + rnd_state = self.get_rnd_state(obs) + self.state_normalizer.update(rnd_state) + + """ + Different weight schedules. + """ + + def _constant_weight_schedule(self, step: int, **kwargs): + return self.initial_weight + + def _step_weight_schedule(self, step: int, final_step: int, final_value: float, **kwargs): + return self.initial_weight if step < final_step else final_value + + def _linear_weight_schedule(self, step: int, initial_step: int, final_step: int, final_value: float, **kwargs): + if step < initial_step: + return self.initial_weight + elif step > final_step: + return final_value + else: + return self.initial_weight + (final_value - self.initial_weight) * (step - initial_step) / ( + final_step - initial_step + ) + + +def resolve_rnd_config(alg_cfg, obs, obs_groups, env): + """Resolve the RND configuration. + + Args: + alg_cfg: The algorithm configuration dictionary. + obs: The observation dictionary. + obs_groups: The observation groups dictionary. + env: The environment. + + Returns: + The resolved algorithm configuration dictionary. + """ + # resolve dimension of rnd gated state + if "rnd_cfg" in alg_cfg and alg_cfg["rnd_cfg"] is not None: + # get dimension of rnd gated state + num_rnd_state = 0 + for obs_group in obs_groups["rnd_state"]: + assert len(obs[obs_group].shape) == 2, "The RND module only supports 1D observations." + num_rnd_state += obs[obs_group].shape[-1] + # add rnd gated state to config + alg_cfg["rnd_cfg"]["num_states"] = num_rnd_state + alg_cfg["rnd_cfg"]["obs_groups"] = obs_groups + # scale down the rnd weight with timestep + alg_cfg["rnd_cfg"]["weight"] *= env.unwrapped.step_dt + return alg_cfg diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py new file mode 100644 index 0000000..6bf1380 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py @@ -0,0 +1,206 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +from torch.distributions import Normal + +from rsl_rl.networks import MLP, EmpiricalNormalization + + +class StudentTeacher(nn.Module): + is_recurrent = False + + def __init__( + self, + obs, + obs_groups, + num_actions, + student_obs_normalization=False, + teacher_obs_normalization=False, + student_hidden_dims=[256, 256, 256], + teacher_hidden_dims=[256, 256, 256], + activation="elu", + init_noise_std=0.1, + noise_std_type: str = "scalar", + **kwargs, + ): + if kwargs: + print( + "StudentTeacher.__init__ got unexpected arguments, which will be ignored: " + + str([key for key in kwargs.keys()]) + ) + super().__init__() + + self.loaded_teacher = False # indicates if teacher has been loaded + + # get the observation dimensions + self.obs_groups = obs_groups + num_student_obs = 0 + for obs_group in obs_groups["policy"]: + assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." + num_student_obs += obs[obs_group].shape[-1] + num_teacher_obs = 0 + for obs_group in obs_groups["teacher"]: + assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." + num_teacher_obs += obs[obs_group].shape[-1] + + # student + self.student = MLP(num_student_obs, num_actions, student_hidden_dims, activation) + + # student observation normalization + self.student_obs_normalization = student_obs_normalization + if student_obs_normalization: + self.student_obs_normalizer = EmpiricalNormalization(num_student_obs) + else: + self.student_obs_normalizer = torch.nn.Identity() + + print(f"Student MLP: {self.student}") + + # teacher + self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation) + self.teacher.eval() + + # teacher observation normalization + self.teacher_obs_normalization = teacher_obs_normalization + if teacher_obs_normalization: + self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs) + else: + self.teacher_obs_normalizer = torch.nn.Identity() + + print(f"Teacher MLP: {self.teacher}") + + # action noise + self.noise_std_type = noise_std_type + if self.noise_std_type == "scalar": + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + elif self.noise_std_type == "log": + self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + + # action distribution (populated in update_distribution) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args(False) + + def reset(self, dones=None, hidden_states=None): + pass + + def forward(self): + raise NotImplementedError + + @property + def action_mean(self): + return self.distribution.mean + + @property + def action_std(self): + return self.distribution.stddev + + @property + def entropy(self): + return self.distribution.entropy().sum(dim=-1) + + def update_distribution(self, obs): + # compute mean + mean = self.student(obs) + # compute standard deviation + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + # create distribution + self.distribution = Normal(mean, std) + + def act(self, obs): + obs = self.get_student_obs(obs) + obs = self.student_obs_normalizer(obs) + self.update_distribution(obs) + return self.distribution.sample() + + def act_inference(self, obs): + obs = self.get_student_obs(obs) + obs = self.student_obs_normalizer(obs) + return self.student(obs) + + def evaluate(self, obs): + obs = self.get_teacher_obs(obs) + obs = self.teacher_obs_normalizer(obs) + with torch.no_grad(): + return self.teacher(obs) + + def get_student_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["policy"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_teacher_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["teacher"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_hidden_states(self): + return None + + def detach_hidden_states(self, dones=None): + pass + + def train(self, mode=True): + super().train(mode) + # make sure teacher is in eval mode + self.teacher.eval() + self.teacher_obs_normalizer.eval() + + def update_normalization(self, obs): + if self.student_obs_normalization: + student_obs = self.get_student_obs(obs) + self.student_obs_normalizer.update(student_obs) + + def load_state_dict(self, state_dict, strict=True): + """Load the parameters of the student and teacher networks. + + Args: + state_dict (dict): State dictionary of the model. + strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this + module's state_dict() function. + + Returns: + bool: Whether this training resumes a previous training. This flag is used by the `load()` function of + `OnPolicyRunner` to determine how to load further parameters. + """ + + # check if state_dict contains teacher and student or just teacher parameters + if any("actor" in key for key in state_dict.keys()): # loading parameters from rl training + # rename keys to match teacher and remove critic parameters + teacher_state_dict = {} + teacher_obs_normalizer_state_dict = {} + for key, value in state_dict.items(): + if "actor." in key: + teacher_state_dict[key.replace("actor.", "")] = value + if "actor_obs_normalizer." in key: + teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value + self.teacher.load_state_dict(teacher_state_dict, strict=strict) + self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict) + # set flag for successfully loading the parameters + self.loaded_teacher = True + self.teacher.eval() + self.teacher_obs_normalizer.eval() + return False # training does not resume + elif any("student" in key for key in state_dict.keys()): # loading parameters from distillation training + super().load_state_dict(state_dict, strict=strict) + # set flag for successfully loading the parameters + self.loaded_teacher = True + self.teacher.eval() + self.teacher_obs_normalizer.eval() + return True # training resumes + else: + raise ValueError("state_dict does not contain student or teacher parameters") diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py new file mode 100644 index 0000000..964a2dc --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py @@ -0,0 +1,249 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +import warnings +from torch.distributions import Normal + +from rsl_rl.networks import MLP, EmpiricalNormalization, Memory + + +class StudentTeacherRecurrent(nn.Module): + is_recurrent = True + + def __init__( + self, + obs, + obs_groups, + num_actions, + student_obs_normalization=False, + teacher_obs_normalization=False, + student_hidden_dims=[256, 256, 256], + teacher_hidden_dims=[256, 256, 256], + activation="elu", + init_noise_std=0.1, + noise_std_type: str = "scalar", + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=1, + teacher_recurrent=False, + **kwargs, + ): + if "rnn_hidden_size" in kwargs: + warnings.warn( + "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. " + "Please use `rnn_hidden_dim` instead.", + DeprecationWarning, + ) + if rnn_hidden_dim == 256: # Only override if the new argument is at its default + rnn_hidden_dim = kwargs.pop("rnn_hidden_size") + if kwargs: + print( + "StudentTeacherRecurrent.__init__ got unexpected arguments, which will be ignored: " + + str(kwargs.keys()), + ) + super().__init__() + + self.loaded_teacher = False # indicates if teacher has been loaded + self.teacher_recurrent = teacher_recurrent # indicates if teacher is recurrent too + + # get the observation dimensions + self.obs_groups = obs_groups + num_student_obs = 0 + for obs_group in obs_groups["policy"]: + assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." + num_student_obs += obs[obs_group].shape[-1] + num_teacher_obs = 0 + for obs_group in obs_groups["teacher"]: + assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." + num_teacher_obs += obs[obs_group].shape[-1] + + # student + self.memory_s = Memory(num_student_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) + self.student = MLP(rnn_hidden_dim, num_actions, student_hidden_dims, activation) + + # student observation normalization + self.student_obs_normalization = student_obs_normalization + if student_obs_normalization: + self.student_obs_normalizer = EmpiricalNormalization(num_student_obs) + else: + self.student_obs_normalizer = torch.nn.Identity() + + print(f"Student RNN: {self.memory_s}") + print(f"Student MLP: {self.student}") + + # teacher + if self.teacher_recurrent: + self.memory_t = Memory( + num_teacher_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim + ) + num_teacher_obs = rnn_hidden_dim + self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation) + + # teacher observation normalization + self.teacher_obs_normalization = teacher_obs_normalization + if teacher_obs_normalization: + self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs) + else: + self.teacher_obs_normalizer = torch.nn.Identity() + + if self.teacher_recurrent: + print(f"Teacher RNN: {self.memory_t}") + print(f"Teacher MLP: {self.teacher}") + + # action noise + self.noise_std_type = noise_std_type + if self.noise_std_type == "scalar": + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + elif self.noise_std_type == "log": + self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + + # action distribution (populated in update_distribution) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args(False) + + def reset(self, dones=None, hidden_states=None): + if hidden_states is None: + hidden_states = (None, None) + self.memory_s.reset(dones, hidden_states[0]) + if self.teacher_recurrent: + self.memory_t.reset(dones, hidden_states[1]) + + def forward(self): + raise NotImplementedError + + @property + def action_mean(self): + return self.distribution.mean + + @property + def action_std(self): + return self.distribution.stddev + + @property + def entropy(self): + return self.distribution.entropy().sum(dim=-1) + + def update_distribution(self, obs): + # compute mean + mean = self.student(obs) + # compute standard deviation + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + # create distribution + self.distribution = Normal(mean, std) + + def act(self, obs): + obs = self.get_student_obs(obs) + obs = self.student_obs_normalizer(obs) + out_mem = self.memory_s(obs).squeeze(0) + self.update_distribution(out_mem) + return self.distribution.sample() + + def act_inference(self, obs): + obs = self.get_student_obs(obs) + obs = self.student_obs_normalizer(obs) + out_mem = self.memory_s(obs).squeeze(0) + return self.student(out_mem) + + def evaluate(self, obs): + obs = self.get_teacher_obs(obs) + obs = self.teacher_obs_normalizer(obs) + with torch.no_grad(): + if self.teacher_recurrent: + self.memory_t.eval() + obs = self.memory_t(obs).squeeze(0) + return self.teacher(obs) + + def get_student_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["policy"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_teacher_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["teacher"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_hidden_states(self): + if self.teacher_recurrent: + return self.memory_s.hidden_states, self.memory_t.hidden_states + else: + return self.memory_s.hidden_states, None + + def detach_hidden_states(self, dones=None): + self.memory_s.detach_hidden_states(dones) + if self.teacher_recurrent: + self.memory_t.detach_hidden_states(dones) + + def train(self, mode=True): + super().train(mode) + # make sure teacher is in eval mode + self.teacher.eval() + self.teacher_obs_normalizer.eval() + + def update_normalization(self, obs): + if self.student_obs_normalization: + student_obs = self.get_student_obs(obs) + self.student_obs_normalizer.update(student_obs) + + def load_state_dict(self, state_dict, strict=True): + """Load the parameters of the student and teacher networks. + + Args: + state_dict (dict): State dictionary of the model. + strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this + module's state_dict() function. + + Returns: + bool: Whether this training resumes a previous training. This flag is used by the `load()` function of + `OnPolicyRunner` to determine how to load further parameters. + """ + + # check if state_dict contains teacher and student or just teacher parameters + if any("actor" in key for key in state_dict.keys()): # loading parameters from rl training + # rename keys to match teacher and remove critic parameters + teacher_state_dict = {} + teacher_obs_normalizer_state_dict = {} + for key, value in state_dict.items(): + if "actor." in key: + teacher_state_dict[key.replace("actor.", "")] = value + if "actor_obs_normalizer." in key: + teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value + self.teacher.load_state_dict(teacher_state_dict, strict=strict) + self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict) + # also load recurrent memory if teacher is recurrent + if self.teacher_recurrent: + memory_t_state_dict = {} + for key, value in state_dict.items(): + if "memory_a." in key: + memory_t_state_dict[key.replace("memory_a.", "")] = value + self.memory_t.load_state_dict(memory_t_state_dict, strict=strict) + # set flag for successfully loading the parameters + self.loaded_teacher = True + self.teacher.eval() + self.teacher_obs_normalizer.eval() + return False # training does not resume + elif any("student" in key for key in state_dict.keys()): # loading parameters from distillation training + super().load_state_dict(state_dict, strict=strict) + # set flag for successfully loading the parameters + self.loaded_teacher = True + self.teacher.eval() + self.teacher_obs_normalizer.eval() + return True # training resumes + else: + raise ValueError("state_dict does not contain student or teacher parameters") diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/symmetry.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/symmetry.py new file mode 100644 index 0000000..b017515 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/symmetry.py @@ -0,0 +1,24 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + + +def resolve_symmetry_config(alg_cfg, env): + """Resolve the symmetry configuration. + + Args: + alg_cfg: The algorithm configuration dictionary. + env: The environment. + + Returns: + The resolved algorithm configuration dictionary. + """ + + # if using symmetry then pass the environment config object + if "symmetry_cfg" in alg_cfg and alg_cfg["symmetry_cfg"] is not None: + # this is used by the symmetry function for handling different observation terms + alg_cfg["symmetry_cfg"]["_env"] = env + return alg_cfg diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/__init__.py new file mode 100644 index 0000000..c18f487 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Definitions for components of modules.""" + +from .memory import Memory +from .mlp import MLP +from .normalization import EmpiricalDiscountedVariationNormalization, EmpiricalNormalization diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/memory.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/memory.py new file mode 100644 index 0000000..7577357 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/memory.py @@ -0,0 +1,70 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch.nn as nn + +from rsl_rl.utils import unpad_trajectories + + +class Memory(nn.Module): + """Memory module for recurrent networks. + + This module is used to store the hidden states of the policy. + Currently only supports GRU and LSTM. + """ + + def __init__(self, input_size, type="lstm", num_layers=1, hidden_size=256): + super().__init__() + # RNN + rnn_cls = nn.GRU if type.lower() == "gru" else nn.LSTM + self.rnn = rnn_cls(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers) + self.hidden_states = None + + def forward(self, input, masks=None, hidden_states=None): + batch_mode = masks is not None + if batch_mode: + # batch mode: needs saved hidden states + if hidden_states is None: + raise ValueError("Hidden states not passed to memory module during policy update") + out, _ = self.rnn(input, hidden_states) + out = unpad_trajectories(out, masks) + else: + # inference/distillation mode: uses hidden states of last step + out, self.hidden_states = self.rnn(input.unsqueeze(0), self.hidden_states) + return out + + def reset(self, dones=None, hidden_states=None): + if dones is None: # reset all hidden states + if hidden_states is None: + self.hidden_states = None + else: + self.hidden_states = hidden_states + elif self.hidden_states is not None: # reset hidden states of done environments + if hidden_states is None: + if isinstance(self.hidden_states, tuple): # tuple in case of LSTM + for hidden_state in self.hidden_states: + hidden_state[..., dones == 1, :] = 0.0 + else: + self.hidden_states[..., dones == 1, :] = 0.0 + else: + NotImplementedError( + "Resetting hidden states of done environments with custom hidden states is not implemented" + ) + + def detach_hidden_states(self, dones=None): + if self.hidden_states is not None: + if dones is None: # detach all hidden states + if isinstance(self.hidden_states, tuple): # tuple in case of LSTM + self.hidden_states = tuple(hidden_state.detach() for hidden_state in self.hidden_states) + else: + self.hidden_states = self.hidden_states.detach() + else: # detach hidden states of done environments + if isinstance(self.hidden_states, tuple): # tuple in case of LSTM + for hidden_state in self.hidden_states: + hidden_state[..., dones == 1, :] = hidden_state[..., dones == 1, :].detach() + else: + self.hidden_states[..., dones == 1, :] = self.hidden_states[..., dones == 1, :].detach() diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/mlp.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/mlp.py new file mode 100644 index 0000000..e91574e --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/mlp.py @@ -0,0 +1,120 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +from functools import reduce + +from rsl_rl.utils import resolve_nn_activation + + +class MLP(nn.Sequential): + """Multi-layer perceptron. + + The MLP network is a sequence of linear layers and activation functions. The + last layer is a linear layer that outputs the desired dimension unless the + last activation function is specified. + + It provides additional conveniences: + + - If the hidden dimensions have a value of ``-1``, the dimension is inferred + from the input dimension. + - If the output dimension is a tuple, the output is reshaped to the desired + shape. + + """ + + def __init__( + self, + input_dim: int, + output_dim: int | tuple[int] | list[int], + hidden_dims: tuple[int] | list[int], + activation: str = "elu", + last_activation: str | None = None, + ): + """Initialize the MLP. + + Args: + input_dim: Dimension of the input. + output_dim: Dimension of the output. + hidden_dims: Dimensions of the hidden layers. A value of ``-1`` indicates + that the dimension should be inferred from the input dimension. + activation: Activation function. Defaults to "elu". + last_activation: Activation function of the last layer. Defaults to None, + in which case the last layer is linear. + """ + super().__init__() + + # resolve activation functions + activation_mod = resolve_nn_activation(activation) + last_activation_mod = resolve_nn_activation(last_activation) if last_activation is not None else None + # resolve number of hidden dims if they are -1 + hidden_dims_processed = [input_dim if dim == -1 else dim for dim in hidden_dims] + + # create layers sequentially + layers = [] + layers.append(nn.Linear(input_dim, hidden_dims_processed[0])) + layers.append(activation_mod) + + for layer_index in range(len(hidden_dims_processed) - 1): + layers.append(nn.Linear(hidden_dims_processed[layer_index], hidden_dims_processed[layer_index + 1])) + layers.append(activation_mod) + + # add last layer + if isinstance(output_dim, int): + layers.append(nn.Linear(hidden_dims_processed[-1], output_dim)) + else: + # compute the total output dimension + total_out_dim = reduce(lambda x, y: x * y, output_dim) + # add a layer to reshape the output to the desired shape + layers.append(nn.Linear(hidden_dims_processed[-1], total_out_dim)) + layers.append(nn.Unflatten(dim=-1, unflattened_size=output_dim)) + + # add last activation function if specified + if last_activation_mod is not None: + layers.append(last_activation_mod) + + # register the layers + for idx, layer in enumerate(layers): + self.add_module(f"{idx}", layer) + + def init_weights(self, scales: float | tuple[float]): + """Initialize the weights of the MLP. + + Args: + scales: Scale factor for the weights. + """ + + def get_scale(idx) -> float: + """Get the scale factor for the weights of the MLP. + + Args: + idx: Index of the layer. + """ + return scales[idx] if isinstance(scales, (list, tuple)) else scales + + # initialize the weights + for idx, module in enumerate(self): + if isinstance(module, nn.Linear): + nn.init.orthogonal_(module.weight, gain=get_scale(idx)) + nn.init.zeros_(module.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass of the MLP. + + Args: + x: Input tensor. + """ + for layer in self: + x = layer(x) + return x + + def reset(self, dones=None, hidden_states=None): + pass + + def detach_hidden_states(self, dones=None): + pass diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/normalization.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/normalization.py new file mode 100644 index 0000000..5fd9692 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/normalization.py @@ -0,0 +1,130 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Copyright (c) 2020 Preferred Networks, Inc. + +from __future__ import annotations + +import torch +from torch import nn + + +class EmpiricalNormalization(nn.Module): + """Normalize mean and variance of values based on empirical values.""" + + def __init__(self, shape, eps=1e-2, until=None): + """Initialize EmpiricalNormalization module. + + Args: + shape (int or tuple of int): Shape of input values except batch axis. + eps (float): Small value for stability. + until (int or None): If this arg is specified, the module learns input values until the sum of batch sizes + exceeds it. + + Note: The normalization parameters are computed over the whole batch, not for each environment separately. + """ + super().__init__() + self.eps = eps + self.until = until + self.register_buffer("_mean", torch.zeros(shape).unsqueeze(0)) + self.register_buffer("_var", torch.ones(shape).unsqueeze(0)) + self.register_buffer("_std", torch.ones(shape).unsqueeze(0)) + self.register_buffer("count", torch.tensor(0, dtype=torch.long)) + + @property + def mean(self): + return self._mean.squeeze(0).clone() + + @property + def std(self): + return self._std.squeeze(0).clone() + + def forward(self, x): + """Normalize mean and variance of values based on empirical values.""" + + return (x - self._mean) / (self._std + self.eps) + + @torch.jit.unused + def update(self, x): + """Learn input values without computing the output values of them""" + + if not self.training: + return + if self.until is not None and self.count >= self.until: + return + + count_x = x.shape[0] + self.count += count_x + rate = count_x / self.count + var_x = torch.var(x, dim=0, unbiased=False, keepdim=True) + mean_x = torch.mean(x, dim=0, keepdim=True) + delta_mean = mean_x - self._mean + self._mean += rate * delta_mean + self._var += rate * (var_x - self._var + delta_mean * (mean_x - self._mean)) + self._std = torch.sqrt(self._var) + + @torch.jit.unused + def inverse(self, y): + """De-normalize values based on empirical values.""" + + return y * (self._std + self.eps) + self._mean + + +class EmpiricalDiscountedVariationNormalization(nn.Module): + """Reward normalization from Pathak's large scale study on PPO. + + Reward normalization. Since the reward function is non-stationary, it is useful to normalize + the scale of the rewards so that the value function can learn quickly. We did this by dividing + the rewards by a running estimate of the standard deviation of the sum of discounted rewards. + """ + + def __init__(self, shape, eps=1e-2, gamma=0.99, until=None): + super().__init__() + + self.emp_norm = EmpiricalNormalization(shape, eps, until) + self.disc_avg = _DiscountedAverage(gamma) + + def forward(self, rew): + if self.training: + # update discounted rewards + avg = self.disc_avg.update(rew) + # update moments from discounted rewards + self.emp_norm.update(avg) + + # normalize rewards with the empirical std + if self.emp_norm._std > 0: + return rew / self.emp_norm._std + else: + return rew + + +""" +Helper class. +""" + + +class _DiscountedAverage: + r"""Discounted average of rewards. + + The discounted average is defined as: + + .. math:: + + \bar{R}_t = \gamma \bar{R}_{t-1} + r_t + + Args: + gamma (float): Discount factor. + """ + + def __init__(self, gamma): + self.avg = None + self.gamma = gamma + + def update(self, rew: torch.Tensor) -> torch.Tensor: + if self.avg is None: + self.avg = rew + else: + self.avg = self.avg * self.gamma + rew + return self.avg diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/__init__.py new file mode 100644 index 0000000..61f1682 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Implementation of runners for environment-agent interaction.""" + +from .on_policy_runner import OnPolicyRunner # isort:skip +from .distillation_runner import DistillationRunner +from .amp_on_policy_runner import AMPOnPolicyRunner + +__all__ = ["OnPolicyRunner", "DistillationRunner", "AMPOnPolicyRunner"] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py new file mode 100644 index 0000000..c0b9b9e --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py @@ -0,0 +1,521 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +import statistics +import time +import torch +import warnings +from collections import deque + +import rsl_rl +from rsl_rl.algorithms import AMP_PPO +from rsl_rl.env import VecEnv +from rsl_rl.modules import ActorCritic, ActorCriticRecurrent,DiscriminatorMulti, resolve_rnd_config, resolve_symmetry_config +from rsl_rl.utils import resolve_obs_groups, store_code_state, Normalizer, G1_AMPLoader + + +class AMPOnPolicyRunner: + """On-policy runner for training and evaluation of actor-critic methods.""" + + def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): + self.cfg = train_cfg + self.alg_cfg = train_cfg["algorithm"] + self.policy_cfg = train_cfg["policy"] + self.device = device + self.env = env + + # check if multi-gpu is enabled + self._configure_multi_gpu() + + # store training configuration + self.num_steps_per_env = self.cfg["num_steps_per_env"] + self.save_interval = self.cfg["save_interval"] + + # query observations from environment for algorithm construction + obs = self.env.get_observations() + default_sets = ["critic"] + if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None: + default_sets.append("rnd_state") + self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets) + + self.amp_data = G1_AMPLoader( + device, + time_between_frames=1/50.0, + preload_transitions=True, + num_preload_transitions=train_cfg["amp_num_preload_transitions"], + motion_files=train_cfg["amp_motion_files"], + num_frames=train_cfg['amp_num_frames'] + ) + + self.amp_observation_dim = self.amp_data.observation_dim if self.cfg["amp_num_obs"] == 0 else self.cfg["amp_num_obs"] + self.amp_num_frames = 0 if self.cfg["amp_num_frames"] == 0 else self.cfg["amp_num_frames"] + self.amp_normalizer = Normalizer(self.amp_observation_dim) + self.discriminator = DiscriminatorMulti( + self.amp_observation_dim, + train_cfg["amp_reward_coef"], + train_cfg["amp_discr_hidden_dims"], + device, + train_cfg["amp_num_frames"], + train_cfg["amp_task_reward_lerp"], + train_cfg['use_lerp'], + ).to(self.device) + + + # create the algorithm + self.alg = self._construct_algorithm(obs) + + # Decide whether to disable logging + # We only log from the process with rank 0 (main process) + self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 + + # Logging + self.log_dir = log_dir + self.writer = None + self.tot_timesteps = 0 + self.tot_time = 0 + self.current_learning_iteration = 0 + self.git_status_repos = [rsl_rl.__file__] + + def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 + # initialize writer + self._prepare_logging_writer() + + # randomize initial episode lengths (for exploration) + if init_at_random_ep_len: + self.env.episode_length_buf = torch.randint_like( + self.env.episode_length_buf, high=int(self.env.max_episode_length) + ) + + # start learning + obs = self.env.get_observations().to(self.device) + self.train_mode() # switch to train mode (for dropout for example) + amp_obs = self.env.get_amp_observations() + amp_obs = amp_obs.to(self.device) + if self.amp_num_frames != 0: + self.amp_obs_frames = torch.zeros(size=(self.env.num_envs, self.amp_num_frames, self.amp_observation_dim), device=self.device) + self.amp_obs_frames = torch.concat((self.amp_obs_frames[:, 1:], amp_obs.unsqueeze(1)), dim=1) + + # Book keeping + ep_infos = [] + rewbuffer = deque(maxlen=100) + lenbuffer = deque(maxlen=100) + step_discrewbuffer = deque(maxlen=100) + + cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_single_step_disc_rew = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + # create buffers for logging extrinsic and intrinsic rewards + if self.alg.rnd: + erewbuffer = deque(maxlen=100) + irewbuffer = deque(maxlen=100) + cur_ereward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_ireward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + # Ensure all parameters are in-synced + if self.is_distributed: + print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") + self.alg.broadcast_parameters() + + # Start training + start_iter = self.current_learning_iteration + tot_iter = start_iter + num_learning_iterations + for it in range(start_iter, tot_iter): + start = time.time() + # Rollout + with torch.inference_mode(): + for _ in range(self.num_steps_per_env): + # Sample actions + actions = self.alg.act(obs,amp_obs) + # Step the environment + obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) + # Move to device + obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) + + next_amp_obs = self.env.get_amp_observations() + next_amp_obs = next_amp_obs.to(self.device) + next_amp_obs_with_term = torch.clone(next_amp_obs) + + reset_env_ids = self.env.reset_env_ids + terminal_amp_states = self.env.get_amp_observations()[reset_env_ids] + next_amp_obs_with_term[reset_env_ids] = terminal_amp_states + self.amp_obs_frames = torch.concat((self.amp_obs_frames[:, 1:], next_amp_obs_with_term.unsqueeze(1)), dim=1) + + amp_reward = torch.zeros(self.env.num_envs, device=obs.device) + + mask = self.env.contact_phase[:, 0] == 1.0 + if mask.any(): + rewards[mask], logit, disc_reward = self.alg.discriminator.predict_amp_reward( + self.amp_obs_frames[mask], rewards[mask], normalizer=self.alg.amp_normalizer + ) + amp_reward[mask] += disc_reward + + # process the step + self.alg.process_env_step(obs, rewards, dones, extras, next_amp_obs_with_term, self.amp_obs_frames) + self.amp_obs_frames[reset_env_ids] = 0 + + amp_obs = torch.clone(next_amp_obs) + # Extract intrinsic rewards (only for logging) + intrinsic_rewards = self.alg.intrinsic_rewards if self.alg.rnd else None + # book keeping + if self.log_dir is not None: + if "episode" in extras: + ep_infos.append(extras["episode"]) + elif "log" in extras: + ep_infos.append(extras["log"]) + # Update rewards + if self.alg.rnd: + cur_ereward_sum += rewards + cur_ireward_sum += intrinsic_rewards # type: ignore + cur_reward_sum += rewards + intrinsic_rewards + else: + cur_reward_sum += rewards + # Update episode length + cur_episode_length += 1 + cur_single_step_disc_rew += amp_reward + # Clear data for completed episodes + # -- common + new_ids = (dones > 0).nonzero(as_tuple=False) + rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) + lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) + cur_reward_sum[new_ids] = 0 + cur_episode_length[new_ids] = 0 + to_extend_disc = (cur_single_step_disc_rew[new_ids] / self.env.max_episode_length_s)[:, 0].cpu().numpy() + step_discrewbuffer.extend(to_extend_disc.tolist()) + cur_single_step_disc_rew[new_ids] = 0 + # -- intrinsic and extrinsic rewards + if self.alg.rnd: + erewbuffer.extend(cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist()) + irewbuffer.extend(cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist()) + cur_ereward_sum[new_ids] = 0 + cur_ireward_sum[new_ids] = 0 + + stop = time.time() + collection_time = stop - start + start = stop + + # compute returns + self.alg.compute_returns(obs) + + # update policy + loss_dict = self.alg.update() + + stop = time.time() + learn_time = stop - start + self.current_learning_iteration = it + # log info + if self.log_dir is not None and not self.disable_logs: + # Log information + self.log(locals()) + # Save model + if it % self.save_interval == 0: + self.save(os.path.join(self.log_dir, f"model_{it}.pt")) + + # Clear episode infos + ep_infos.clear() + # Save code state + if it == start_iter and not self.disable_logs: + # obtain all the diff files + git_file_paths = store_code_state(self.log_dir, self.git_status_repos) + # if possible store them to wandb + if self.logger_type in ["wandb", "neptune"] and git_file_paths: + for path in git_file_paths: + self.writer.save_file(path) + + # Save the final model after training + if self.log_dir is not None and not self.disable_logs: + self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) + + def log(self, locs: dict, width: int = 80, pad: int = 35): + # Compute the collection size + collection_size = self.num_steps_per_env * self.env.num_envs * self.gpu_world_size + # Update total time-steps and time + self.tot_timesteps += collection_size + self.tot_time += locs["collection_time"] + locs["learn_time"] + iteration_time = locs["collection_time"] + locs["learn_time"] + + # -- Episode info + ep_string = "" + if locs["ep_infos"]: + for key in locs["ep_infos"][0]: + infotensor = torch.tensor([], device=self.device) + for ep_info in locs["ep_infos"]: + # handle scalar and zero dimensional tensor infos + if key not in ep_info: + continue + if not isinstance(ep_info[key], torch.Tensor): + ep_info[key] = torch.Tensor([ep_info[key]]) + if len(ep_info[key].shape) == 0: + ep_info[key] = ep_info[key].unsqueeze(0) + infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) + value = torch.mean(infotensor) + # log to logger and terminal + if "/" in key: + self.writer.add_scalar(key, value, locs["it"]) + ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + else: + self.writer.add_scalar("Episode/" + key, value, locs["it"]) + ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" + + mean_std = self.alg.policy.action_std.mean() + fps = int(collection_size / (locs["collection_time"] + locs["learn_time"])) + + # -- Losses + for key, value in locs["loss_dict"].items(): + self.writer.add_scalar(f"Loss/{key}", value, locs["it"]) + self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) + + # -- Policy + self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) + + # -- Performance + self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) + self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) + self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) + + # -- Training + if len(locs["rewbuffer"]) > 0: + # separate logging for intrinsic and extrinsic rewards + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(locs["erewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(locs["irewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/weight", self.alg.rnd.weight, locs["it"]) + # everything else + self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) + self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) + self.writer.add_scalar('Train/mean_step_disc_reward', statistics.mean(locs['step_discrewbuffer']), locs['it']) + if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging + self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) + self.writer.add_scalar( + "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time + ) + + str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " + + if len(locs["rewbuffer"]) > 0: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + f"""{'Step disc reward:':>{pad}} {statistics.mean(locs['step_discrewbuffer']):.2f}\n""" + ) + # -- Losses + for key, value in locs["loss_dict"].items(): + log_string += f"""{f'Mean {key} loss:':>{pad}} {value:.4f}\n""" + # -- Rewards + if hasattr(self.alg, "rnd") and self.alg.rnd: + log_string += ( + f"""{'Mean extrinsic reward:':>{pad}} {statistics.mean(locs['erewbuffer']):.2f}\n""" + f"""{'Mean intrinsic reward:':>{pad}} {statistics.mean(locs['irewbuffer']):.2f}\n""" + ) + log_string += f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" + # -- episode info + log_string += f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" + else: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + ) + for key, value in locs["loss_dict"].items(): + log_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + + log_string += ep_string + log_string += ( + f"""{'-' * width}\n""" + f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" + f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" + f"""{'Time elapsed:':>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n""" + f"""{'ETA:':>{pad}} {time.strftime( + "%H:%M:%S", + time.gmtime( + self.tot_time / (locs['it'] - locs['start_iter'] + 1) + * (locs['start_iter'] + locs['num_learning_iterations'] - locs['it']) + ) + )}\n""" + ) + print(log_string) + + def save(self, path: str, infos=None): + # -- Save model + saved_dict = { + "model_state_dict": self.alg.policy.state_dict(), + "optimizer_state_dict": self.alg.optimizer.state_dict(), + "iter": self.current_learning_iteration, + "infos": infos, + } + # -- Save RND model if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict() + saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict() + torch.save(saved_dict, path) + + # upload model to external logging service + if self.logger_type in ["neptune", "wandb"] and not self.disable_logs: + self.writer.save_model(path, self.current_learning_iteration) + + def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None): + loaded_dict = torch.load(path, weights_only=False, map_location=map_location) + # -- Load model + resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"]) + # -- Load RND model if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"]) + # -- load optimizer if used + if load_optimizer and resumed_training: + # -- algorithm optimizer + self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) + # -- RND optimizer if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"]) + # -- load current learning iteration + if resumed_training: + self.current_learning_iteration = loaded_dict["iter"] + return loaded_dict["infos"] + + def get_inference_policy(self, device=None): + self.eval_mode() # switch to evaluation mode (dropout for example) + if device is not None: + self.alg.policy.to(device) + return self.alg.policy.act_inference + + def train_mode(self): + # -- PPO + self.alg.policy.train() + self.alg.discriminator.train() + # -- RND + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.train() + + def eval_mode(self): + # -- PPO + self.alg.policy.eval() + self.alg.discriminator.eval() + # -- RND + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.eval() + + def add_git_repo_to_log(self, repo_file_path): + self.git_status_repos.append(repo_file_path) + + """ + Helper functions. + """ + + def _configure_multi_gpu(self): + """Configure multi-gpu training.""" + # check if distributed training is enabled + self.gpu_world_size = int(os.getenv("WORLD_SIZE", "1")) + self.is_distributed = self.gpu_world_size > 1 + + # if not distributed training, set local and global rank to 0 and return + if not self.is_distributed: + self.gpu_local_rank = 0 + self.gpu_global_rank = 0 + self.multi_gpu_cfg = None + return + + # get rank and world size + self.gpu_local_rank = int(os.getenv("LOCAL_RANK", "0")) + self.gpu_global_rank = int(os.getenv("RANK", "0")) + + # make a configuration dictionary + self.multi_gpu_cfg = { + "global_rank": self.gpu_global_rank, # rank of the main process + "local_rank": self.gpu_local_rank, # rank of the current process + "world_size": self.gpu_world_size, # total number of processes + } + + # check if user has device specified for local rank + if self.device != f"cuda:{self.gpu_local_rank}": + raise ValueError( + f"Device '{self.device}' does not match expected device for local rank '{self.gpu_local_rank}'." + ) + # validate multi-gpu configuration + if self.gpu_local_rank >= self.gpu_world_size: + raise ValueError( + f"Local rank '{self.gpu_local_rank}' is greater than or equal to world size '{self.gpu_world_size}'." + ) + if self.gpu_global_rank >= self.gpu_world_size: + raise ValueError( + f"Global rank '{self.gpu_global_rank}' is greater than or equal to world size '{self.gpu_world_size}'." + ) + + # initialize torch distributed + torch.distributed.init_process_group(backend="nccl", rank=self.gpu_global_rank, world_size=self.gpu_world_size) + # set device to the local rank + torch.cuda.set_device(self.gpu_local_rank) + + def _construct_algorithm(self, obs) -> AMP_PPO: + """Construct the actor-critic algorithm.""" + # resolve RND config + self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env) + + # resolve symmetry config + self.alg_cfg = resolve_symmetry_config(self.alg_cfg, self.env) + + # resolve deprecated normalization config + if self.cfg.get("empirical_normalization") is not None: + warnings.warn( + "The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and " + "`critic_obs_normalization` as part of the `policy` configuration instead.", + DeprecationWarning, + ) + if self.policy_cfg.get("actor_obs_normalization") is None: + self.policy_cfg["actor_obs_normalization"] = self.cfg["empirical_normalization"] + if self.policy_cfg.get("critic_obs_normalization") is None: + self.policy_cfg["critic_obs_normalization"] = self.cfg["empirical_normalization"] + + # initialize the actor-critic + actor_critic_class = eval(self.policy_cfg.pop("class_name")) + actor_critic: ActorCritic | ActorCriticRecurrent = actor_critic_class( + obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg + ).to(self.device) + + # initialize the algorithm + alg_class = eval(self.alg_cfg.pop("class_name")) + + alg: AMP_PPO = alg_class(actor_critic, self.discriminator, self.amp_data, self.amp_normalizer, self.amp_num_frames, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg) + + # initialize the storage + alg.init_storage( + "rl", + self.env.num_envs, + self.num_steps_per_env, + obs, + [self.env.num_actions], + ) + + return alg + + def _prepare_logging_writer(self): + """Prepares the logging writers.""" + if self.log_dir is not None and self.writer is None and not self.disable_logs: + # Launch either Tensorboard or Neptune & Tensorboard summary writer(s), default: Tensorboard. + self.logger_type = self.cfg.get("logger", "tensorboard") + self.logger_type = self.logger_type.lower() + + if self.logger_type == "neptune": + from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter + + self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "wandb": + from rsl_rl.utils.wandb_utils import WandbSummaryWriter + + self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "tensorboard": + from torch.utils.tensorboard import SummaryWriter + + self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) + else: + raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.") diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py new file mode 100644 index 0000000..9cc6a8b --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py @@ -0,0 +1,179 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +import time +import torch +from collections import deque + +import rsl_rl +from rsl_rl.algorithms import Distillation +from rsl_rl.env import VecEnv +from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent +from rsl_rl.runners import OnPolicyRunner +from rsl_rl.utils import resolve_obs_groups, store_code_state + + +class DistillationRunner(OnPolicyRunner): + """On-policy runner for training and evaluation of teacher-student training.""" + + def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): + self.cfg = train_cfg + self.alg_cfg = train_cfg["algorithm"] + self.policy_cfg = train_cfg["policy"] + self.device = device + self.env = env + + # check if multi-gpu is enabled + self._configure_multi_gpu() + + # store training configuration + self.num_steps_per_env = self.cfg["num_steps_per_env"] + self.save_interval = self.cfg["save_interval"] + + # query observations from environment for algorithm construction + obs = self.env.get_observations() + self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets=["teacher"]) + + # create the algorithm + self.alg = self._construct_algorithm(obs) + + # Decide whether to disable logging + # We only log from the process with rank 0 (main process) + self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 + + # Logging + self.log_dir = log_dir + self.writer = None + self.tot_timesteps = 0 + self.tot_time = 0 + self.current_learning_iteration = 0 + self.git_status_repos = [rsl_rl.__file__] + + def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 + # initialize writer + self._prepare_logging_writer() + # check if teacher is loaded + if not self.alg.policy.loaded_teacher: + raise ValueError("Teacher model parameters not loaded. Please load a teacher model to distill.") + + # randomize initial episode lengths (for exploration) + if init_at_random_ep_len: + self.env.episode_length_buf = torch.randint_like( + self.env.episode_length_buf, high=int(self.env.max_episode_length) + ) + + # start learning + obs = self.env.get_observations().to(self.device) + self.train_mode() # switch to train mode (for dropout for example) + + # Book keeping + ep_infos = [] + rewbuffer = deque(maxlen=100) + lenbuffer = deque(maxlen=100) + cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + # Ensure all parameters are in-synced + if self.is_distributed: + print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") + self.alg.broadcast_parameters() + + # Start training + start_iter = self.current_learning_iteration + tot_iter = start_iter + num_learning_iterations + for it in range(start_iter, tot_iter): + start = time.time() + # Rollout + with torch.inference_mode(): + for _ in range(self.num_steps_per_env): + # Sample actions + actions = self.alg.act(obs) + # Step the environment + obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) + # Move to device + obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) + # process the step + self.alg.process_env_step(obs, rewards, dones, extras) + # book keeping + if self.log_dir is not None: + if "episode" in extras: + ep_infos.append(extras["episode"]) + elif "log" in extras: + ep_infos.append(extras["log"]) + # Update rewards + cur_reward_sum += rewards + # Update episode length + cur_episode_length += 1 + # Clear data for completed episodes + new_ids = (dones > 0).nonzero(as_tuple=False) + rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) + lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) + cur_reward_sum[new_ids] = 0 + cur_episode_length[new_ids] = 0 + + stop = time.time() + collection_time = stop - start + start = stop + + # update policy + loss_dict = self.alg.update() + + stop = time.time() + learn_time = stop - start + self.current_learning_iteration = it + # log info + if self.log_dir is not None and not self.disable_logs: + # Log information + self.log(locals()) + # Save model + if it % self.save_interval == 0: + self.save(os.path.join(self.log_dir, f"model_{it}.pt")) + + # Clear episode infos + ep_infos.clear() + # Save code state + if it == start_iter and not self.disable_logs: + # obtain all the diff files + git_file_paths = store_code_state(self.log_dir, self.git_status_repos) + # if possible store them to wandb + if self.logger_type in ["wandb", "neptune"] and git_file_paths: + for path in git_file_paths: + self.writer.save_file(path) + + # Save the final model after training + if self.log_dir is not None and not self.disable_logs: + self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) + + """ + Helper methods. + """ + + def _construct_algorithm(self, obs) -> Distillation: + """Construct the distillation algorithm.""" + # initialize the actor-critic + student_teacher_class = eval(self.policy_cfg.pop("class_name")) + student_teacher: StudentTeacher | StudentTeacherRecurrent = student_teacher_class( + obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg + ).to(self.device) + + # initialize the algorithm + alg_class = eval(self.alg_cfg.pop("class_name")) + alg: Distillation = alg_class( + student_teacher, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg + ) + + # initialize the storage + alg.init_storage( + "distillation", + self.env.num_envs, + self.num_steps_per_env, + obs, + [self.env.num_actions], + ) + + return alg diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py new file mode 100644 index 0000000..36f11f3 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py @@ -0,0 +1,460 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +import statistics +import time +import torch +import warnings +from collections import deque + +import rsl_rl +from rsl_rl.algorithms import PPO +from rsl_rl.env import VecEnv +from rsl_rl.modules import ActorCritic, ActorCriticRecurrent, resolve_rnd_config, resolve_symmetry_config +from rsl_rl.utils import resolve_obs_groups, store_code_state + + +class OnPolicyRunner: + """On-policy runner for training and evaluation of actor-critic methods.""" + + def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): + self.cfg = train_cfg + self.alg_cfg = train_cfg["algorithm"] + self.policy_cfg = train_cfg["policy"] + self.device = device + self.env = env + + # check if multi-gpu is enabled + self._configure_multi_gpu() + + # store training configuration + self.num_steps_per_env = self.cfg["num_steps_per_env"] + self.save_interval = self.cfg["save_interval"] + + # query observations from environment for algorithm construction + obs = self.env.get_observations() + default_sets = ["critic"] + if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None: + default_sets.append("rnd_state") + self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets) + + # create the algorithm + self.alg = self._construct_algorithm(obs) + + # Decide whether to disable logging + # We only log from the process with rank 0 (main process) + self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 + + # Logging + self.log_dir = log_dir + self.writer = None + self.tot_timesteps = 0 + self.tot_time = 0 + self.current_learning_iteration = 0 + self.git_status_repos = [rsl_rl.__file__] + + def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 + # initialize writer + self._prepare_logging_writer() + + # randomize initial episode lengths (for exploration) + if init_at_random_ep_len: + self.env.episode_length_buf = torch.randint_like( + self.env.episode_length_buf, high=int(self.env.max_episode_length) + ) + + # start learning + obs = self.env.get_observations().to(self.device) + self.train_mode() # switch to train mode (for dropout for example) + + # Book keeping + ep_infos = [] + rewbuffer = deque(maxlen=100) + lenbuffer = deque(maxlen=100) + cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + # create buffers for logging extrinsic and intrinsic rewards + if self.alg.rnd: + erewbuffer = deque(maxlen=100) + irewbuffer = deque(maxlen=100) + cur_ereward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_ireward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + # Ensure all parameters are in-synced + if self.is_distributed: + print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") + self.alg.broadcast_parameters() + + # Start training + start_iter = self.current_learning_iteration + tot_iter = start_iter + num_learning_iterations + for it in range(start_iter, tot_iter): + start = time.time() + # Rollout + with torch.inference_mode(): + for _ in range(self.num_steps_per_env): + # Sample actions + actions = self.alg.act(obs) + # Step the environment + obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) + # Move to device + obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) + # process the step + self.alg.process_env_step(obs, rewards, dones, extras) + # Extract intrinsic rewards (only for logging) + intrinsic_rewards = self.alg.intrinsic_rewards if self.alg.rnd else None + # book keeping + if self.log_dir is not None: + if "episode" in extras: + ep_infos.append(extras["episode"]) + elif "log" in extras: + ep_infos.append(extras["log"]) + # Update rewards + if self.alg.rnd: + cur_ereward_sum += rewards + cur_ireward_sum += intrinsic_rewards # type: ignore + cur_reward_sum += rewards + intrinsic_rewards + else: + cur_reward_sum += rewards + # Update episode length + cur_episode_length += 1 + # Clear data for completed episodes + # -- common + new_ids = (dones > 0).nonzero(as_tuple=False) + rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) + lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) + cur_reward_sum[new_ids] = 0 + cur_episode_length[new_ids] = 0 + # -- intrinsic and extrinsic rewards + if self.alg.rnd: + erewbuffer.extend(cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist()) + irewbuffer.extend(cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist()) + cur_ereward_sum[new_ids] = 0 + cur_ireward_sum[new_ids] = 0 + + stop = time.time() + collection_time = stop - start + start = stop + + # compute returns + self.alg.compute_returns(obs) + + # update policy + loss_dict = self.alg.update() + + stop = time.time() + learn_time = stop - start + self.current_learning_iteration = it + # log info + if self.log_dir is not None and not self.disable_logs: + # Log information + self.log(locals()) + # Save model + if it % self.save_interval == 0: + self.save(os.path.join(self.log_dir, f"model_{it}.pt")) + + # Clear episode infos + ep_infos.clear() + # Save code state + if it == start_iter and not self.disable_logs: + # obtain all the diff files + git_file_paths = store_code_state(self.log_dir, self.git_status_repos) + # if possible store them to wandb + if self.logger_type in ["wandb", "neptune"] and git_file_paths: + for path in git_file_paths: + self.writer.save_file(path) + + # Save the final model after training + if self.log_dir is not None and not self.disable_logs: + self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) + + def log(self, locs: dict, width: int = 80, pad: int = 35): + # Compute the collection size + collection_size = self.num_steps_per_env * self.env.num_envs * self.gpu_world_size + # Update total time-steps and time + self.tot_timesteps += collection_size + self.tot_time += locs["collection_time"] + locs["learn_time"] + iteration_time = locs["collection_time"] + locs["learn_time"] + + # -- Episode info + ep_string = "" + if locs["ep_infos"]: + for key in locs["ep_infos"][0]: + infotensor = torch.tensor([], device=self.device) + for ep_info in locs["ep_infos"]: + # handle scalar and zero dimensional tensor infos + if key not in ep_info: + continue + if not isinstance(ep_info[key], torch.Tensor): + ep_info[key] = torch.Tensor([ep_info[key]]) + if len(ep_info[key].shape) == 0: + ep_info[key] = ep_info[key].unsqueeze(0) + infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) + value = torch.mean(infotensor) + # log to logger and terminal + if "/" in key: + self.writer.add_scalar(key, value, locs["it"]) + ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + else: + self.writer.add_scalar("Episode/" + key, value, locs["it"]) + ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" + + mean_std = self.alg.policy.action_std.mean() + fps = int(collection_size / (locs["collection_time"] + locs["learn_time"])) + + # -- Losses + for key, value in locs["loss_dict"].items(): + self.writer.add_scalar(f"Loss/{key}", value, locs["it"]) + self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) + + # -- Policy + self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) + + # -- Performance + self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) + self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) + self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) + + # -- Training + if len(locs["rewbuffer"]) > 0: + # separate logging for intrinsic and extrinsic rewards + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(locs["erewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(locs["irewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/weight", self.alg.rnd.weight, locs["it"]) + # everything else + self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) + self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) + if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging + self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) + self.writer.add_scalar( + "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time + ) + + str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " + + if len(locs["rewbuffer"]) > 0: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + ) + # -- Losses + for key, value in locs["loss_dict"].items(): + log_string += f"""{f'Mean {key} loss:':>{pad}} {value:.4f}\n""" + # -- Rewards + if hasattr(self.alg, "rnd") and self.alg.rnd: + log_string += ( + f"""{'Mean extrinsic reward:':>{pad}} {statistics.mean(locs['erewbuffer']):.2f}\n""" + f"""{'Mean intrinsic reward:':>{pad}} {statistics.mean(locs['irewbuffer']):.2f}\n""" + ) + log_string += f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" + # -- episode info + log_string += f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" + else: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + ) + for key, value in locs["loss_dict"].items(): + log_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + + log_string += ep_string + log_string += ( + f"""{'-' * width}\n""" + f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" + f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" + f"""{'Time elapsed:':>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n""" + f"""{'ETA:':>{pad}} {time.strftime( + "%H:%M:%S", + time.gmtime( + self.tot_time / (locs['it'] - locs['start_iter'] + 1) + * (locs['start_iter'] + locs['num_learning_iterations'] - locs['it']) + ) + )}\n""" + ) + print(log_string) + + def save(self, path: str, infos=None): + # -- Save model + saved_dict = { + "model_state_dict": self.alg.policy.state_dict(), + "optimizer_state_dict": self.alg.optimizer.state_dict(), + "iter": self.current_learning_iteration, + "infos": infos, + } + # -- Save RND model if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict() + saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict() + torch.save(saved_dict, path) + + # upload model to external logging service + if self.logger_type in ["neptune", "wandb"] and not self.disable_logs: + self.writer.save_model(path, self.current_learning_iteration) + + def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None): + loaded_dict = torch.load(path, weights_only=False, map_location=map_location) + # -- Load model + resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"]) + # -- Load RND model if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"]) + # -- load optimizer if used + if load_optimizer and resumed_training: + # -- algorithm optimizer + self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) + # -- RND optimizer if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"]) + # -- load current learning iteration + if resumed_training: + self.current_learning_iteration = loaded_dict["iter"] + return loaded_dict["infos"] + + def get_inference_policy(self, device=None): + self.eval_mode() # switch to evaluation mode (dropout for example) + if device is not None: + self.alg.policy.to(device) + return self.alg.policy.act_inference + + def train_mode(self): + # -- PPO + self.alg.policy.train() + # -- RND + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.train() + + def eval_mode(self): + # -- PPO + self.alg.policy.eval() + # -- RND + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.eval() + + def add_git_repo_to_log(self, repo_file_path): + self.git_status_repos.append(repo_file_path) + + """ + Helper functions. + """ + + def _configure_multi_gpu(self): + """Configure multi-gpu training.""" + # check if distributed training is enabled + self.gpu_world_size = int(os.getenv("WORLD_SIZE", "1")) + self.is_distributed = self.gpu_world_size > 1 + + # if not distributed training, set local and global rank to 0 and return + if not self.is_distributed: + self.gpu_local_rank = 0 + self.gpu_global_rank = 0 + self.multi_gpu_cfg = None + return + + # get rank and world size + self.gpu_local_rank = int(os.getenv("LOCAL_RANK", "0")) + self.gpu_global_rank = int(os.getenv("RANK", "0")) + + # make a configuration dictionary + self.multi_gpu_cfg = { + "global_rank": self.gpu_global_rank, # rank of the main process + "local_rank": self.gpu_local_rank, # rank of the current process + "world_size": self.gpu_world_size, # total number of processes + } + + # check if user has device specified for local rank + if self.device != f"cuda:{self.gpu_local_rank}": + raise ValueError( + f"Device '{self.device}' does not match expected device for local rank '{self.gpu_local_rank}'." + ) + # validate multi-gpu configuration + if self.gpu_local_rank >= self.gpu_world_size: + raise ValueError( + f"Local rank '{self.gpu_local_rank}' is greater than or equal to world size '{self.gpu_world_size}'." + ) + if self.gpu_global_rank >= self.gpu_world_size: + raise ValueError( + f"Global rank '{self.gpu_global_rank}' is greater than or equal to world size '{self.gpu_world_size}'." + ) + + # initialize torch distributed + torch.distributed.init_process_group(backend="nccl", rank=self.gpu_global_rank, world_size=self.gpu_world_size) + # set device to the local rank + torch.cuda.set_device(self.gpu_local_rank) + + def _construct_algorithm(self, obs) -> PPO: + """Construct the actor-critic algorithm.""" + # resolve RND config + self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env) + + # resolve symmetry config + self.alg_cfg = resolve_symmetry_config(self.alg_cfg, self.env) + + # resolve deprecated normalization config + if self.cfg.get("empirical_normalization") is not None: + warnings.warn( + "The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and " + "`critic_obs_normalization` as part of the `policy` configuration instead.", + DeprecationWarning, + ) + if self.policy_cfg.get("actor_obs_normalization") is None: + self.policy_cfg["actor_obs_normalization"] = self.cfg["empirical_normalization"] + if self.policy_cfg.get("critic_obs_normalization") is None: + self.policy_cfg["critic_obs_normalization"] = self.cfg["empirical_normalization"] + + # initialize the actor-critic + actor_critic_class = eval(self.policy_cfg.pop("class_name")) + actor_critic: ActorCritic | ActorCriticRecurrent = actor_critic_class( + obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg + ).to(self.device) + + # initialize the algorithm + alg_class = eval(self.alg_cfg.pop("class_name")) + alg: PPO = alg_class(actor_critic, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg) + + # initialize the storage + alg.init_storage( + "rl", + self.env.num_envs, + self.num_steps_per_env, + obs, + [self.env.num_actions], + ) + + return alg + + def _prepare_logging_writer(self): + """Prepares the logging writers.""" + if self.log_dir is not None and self.writer is None and not self.disable_logs: + # Launch either Tensorboard or Neptune & Tensorboard summary writer(s), default: Tensorboard. + self.logger_type = self.cfg.get("logger", "tensorboard") + self.logger_type = self.logger_type.lower() + + if self.logger_type == "neptune": + from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter + + self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "wandb": + from rsl_rl.utils.wandb_utils import WandbSummaryWriter + + self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "tensorboard": + from torch.utils.tensorboard import SummaryWriter + + self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) + else: + raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.") diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/__init__.py new file mode 100644 index 0000000..1624330 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Implementation of transitions storage for RL-agent.""" + +from .rollout_storage import RolloutStorage +from .replay_buffer_multi import ReplayBufferMulti +__all__ = ["RolloutStorage", "ReplayBufferMulti"] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py new file mode 100644 index 0000000..6462b8e --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py @@ -0,0 +1,38 @@ +import torch +import numpy as np + + +class ReplayBufferMulti: + """Fixed-size buffer to store experience tuples.""" + + def __init__(self, obs_dim, buffer_size, num_amp_frames, device): + """Initialize a ReplayBuffer object. + Arguments: + buffer_size (int): maximum size of buffer + """ + self.states = torch.zeros(buffer_size, num_amp_frames, obs_dim).to(device) + self.num_amp_frames = num_amp_frames + self.buffer_size = buffer_size + self.device = device + + self.step = 0 + self.num_samples = 0 + + def insert(self, states): + """Add new states to memory.""" + num_states = states.shape[0] + start_idx = self.step + end_idx = self.step + num_states + if end_idx > self.buffer_size: + self.states[self.step:self.buffer_size] = states[:self.buffer_size - self.step] + self.states[:end_idx - self.buffer_size] = states[self.buffer_size - self.step:] + else: + self.states[start_idx:end_idx] = states + + self.num_samples = min(self.buffer_size, max(end_idx, self.num_samples)) + self.step = (self.step + num_states) % self.buffer_size + + def feed_forward_generator(self, num_mini_batch, mini_batch_size): + for _ in range(num_mini_batch): + sample_idxs = np.random.choice(self.num_samples, size=mini_batch_size) + yield (self.states[sample_idxs].to(self.device)) diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py new file mode 100644 index 0000000..e9309b3 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py @@ -0,0 +1,260 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from tensordict import TensorDict + +from rsl_rl.utils import split_and_pad_trajectories + + +class RolloutStorage: + class Transition: + def __init__(self): + self.observations = None + self.actions = None + self.privileged_actions = None + self.rewards = None + self.dones = None + self.values = None + self.actions_log_prob = None + self.action_mean = None + self.action_sigma = None + self.hidden_states = None + + def clear(self): + self.__init__() + + def __init__( + self, + training_type, + num_envs, + num_transitions_per_env, + obs, + actions_shape, + device="cpu", + ): + # store inputs + self.training_type = training_type + self.device = device + self.num_transitions_per_env = num_transitions_per_env + self.num_envs = num_envs + self.actions_shape = actions_shape + + # Core + self.observations = TensorDict( + {key: torch.zeros(num_transitions_per_env, *value.shape, device=device) for key, value in obs.items()}, + batch_size=[num_transitions_per_env, num_envs], + device=self.device, + ) + self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte() + + # for distillation + if training_type == "distillation": + self.privileged_actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + + # for reinforcement learning + if training_type == "rl": + self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + + # For RNN networks + self.saved_hidden_states_a = None + self.saved_hidden_states_c = None + + # counter for the number of transitions stored + self.step = 0 + + def add_transitions(self, transition: Transition): + # check if the transition is valid + if self.step >= self.num_transitions_per_env: + raise OverflowError("Rollout buffer overflow! You should call clear() before adding new transitions.") + + # Core + self.observations[self.step].copy_(transition.observations) + self.actions[self.step].copy_(transition.actions) + self.rewards[self.step].copy_(transition.rewards.view(-1, 1)) + self.dones[self.step].copy_(transition.dones.view(-1, 1)) + + # for distillation + if self.training_type == "distillation": + self.privileged_actions[self.step].copy_(transition.privileged_actions) + + # for reinforcement learning + if self.training_type == "rl": + self.values[self.step].copy_(transition.values) + self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1)) + self.mu[self.step].copy_(transition.action_mean) + self.sigma[self.step].copy_(transition.action_sigma) + + # For RNN networks + self._save_hidden_states(transition.hidden_states) + + # increment the counter + self.step += 1 + + def _save_hidden_states(self, hidden_states): + if hidden_states is None or hidden_states == (None, None): + return + # make a tuple out of GRU hidden state sto match the LSTM format + hid_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],) + hid_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],) + # initialize if needed + if self.saved_hidden_states_a is None: + self.saved_hidden_states_a = [ + torch.zeros(self.observations.shape[0], *hid_a[i].shape, device=self.device) for i in range(len(hid_a)) + ] + self.saved_hidden_states_c = [ + torch.zeros(self.observations.shape[0], *hid_c[i].shape, device=self.device) for i in range(len(hid_c)) + ] + # copy the states + for i in range(len(hid_a)): + self.saved_hidden_states_a[i][self.step].copy_(hid_a[i]) + self.saved_hidden_states_c[i][self.step].copy_(hid_c[i]) + + def clear(self): + self.step = 0 + + def compute_returns(self, last_values, gamma, lam, normalize_advantage: bool = True): + advantage = 0 + for step in reversed(range(self.num_transitions_per_env)): + # if we are at the last step, bootstrap the return value + if step == self.num_transitions_per_env - 1: + next_values = last_values + else: + next_values = self.values[step + 1] + # 1 if we are not in a terminal state, 0 otherwise + next_is_not_terminal = 1.0 - self.dones[step].float() + # TD error: r_t + gamma * V(s_{t+1}) - V(s_t) + delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step] + # Advantage: A(s_t, a_t) = delta_t + gamma * lambda * A(s_{t+1}, a_{t+1}) + advantage = delta + next_is_not_terminal * gamma * lam * advantage + # Return: R_t = A(s_t, a_t) + V(s_t) + self.returns[step] = advantage + self.values[step] + + # Compute the advantages + self.advantages = self.returns - self.values + # Normalize the advantages if flag is set + # This is to prevent double normalization (i.e. if per minibatch normalization is used) + if normalize_advantage: + self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8) + + # for distillation + def generator(self): + if self.training_type != "distillation": + raise ValueError("This function is only available for distillation training.") + + for i in range(self.num_transitions_per_env): + yield self.observations[i], self.actions[i], self.privileged_actions[i], self.dones[i] + + # for reinforcement learning with feedforward networks + def mini_batch_generator(self, num_mini_batches, num_epochs=8): + if self.training_type != "rl": + raise ValueError("This function is only available for reinforcement learning training.") + batch_size = self.num_envs * self.num_transitions_per_env + mini_batch_size = batch_size // num_mini_batches + indices = torch.randperm(num_mini_batches * mini_batch_size, requires_grad=False, device=self.device) + + # Core + observations = self.observations.flatten(0, 1) + actions = self.actions.flatten(0, 1) + values = self.values.flatten(0, 1) + returns = self.returns.flatten(0, 1) + + # For PPO + old_actions_log_prob = self.actions_log_prob.flatten(0, 1) + advantages = self.advantages.flatten(0, 1) + old_mu = self.mu.flatten(0, 1) + old_sigma = self.sigma.flatten(0, 1) + + for epoch in range(num_epochs): + for i in range(num_mini_batches): + # Select the indices for the mini-batch + start = i * mini_batch_size + end = (i + 1) * mini_batch_size + batch_idx = indices[start:end] + + # Create the mini-batch + # -- Core + obs_batch = observations[batch_idx] + actions_batch = actions[batch_idx] + + # -- For PPO + target_values_batch = values[batch_idx] + returns_batch = returns[batch_idx] + old_actions_log_prob_batch = old_actions_log_prob[batch_idx] + advantages_batch = advantages[batch_idx] + old_mu_batch = old_mu[batch_idx] + old_sigma_batch = old_sigma[batch_idx] + + # yield the mini-batch + yield obs_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( + None, + None, + ), None + + # for reinfrocement learning with recurrent networks + def recurrent_mini_batch_generator(self, num_mini_batches, num_epochs=8): + if self.training_type != "rl": + raise ValueError("This function is only available for reinforcement learning training.") + padded_obs_trajectories, trajectory_masks = split_and_pad_trajectories(self.observations, self.dones) + + mini_batch_size = self.num_envs // num_mini_batches + for ep in range(num_epochs): + first_traj = 0 + for i in range(num_mini_batches): + start = i * mini_batch_size + stop = (i + 1) * mini_batch_size + + dones = self.dones.squeeze(-1) + last_was_done = torch.zeros_like(dones, dtype=torch.bool) + last_was_done[1:] = dones[:-1] + last_was_done[0] = True + trajectories_batch_size = torch.sum(last_was_done[:, start:stop]) + last_traj = first_traj + trajectories_batch_size + + masks_batch = trajectory_masks[:, first_traj:last_traj] + obs_batch = padded_obs_trajectories[:, first_traj:last_traj] + actions_batch = self.actions[:, start:stop] + old_mu_batch = self.mu[:, start:stop] + old_sigma_batch = self.sigma[:, start:stop] + returns_batch = self.returns[:, start:stop] + advantages_batch = self.advantages[:, start:stop] + values_batch = self.values[:, start:stop] + old_actions_log_prob_batch = self.actions_log_prob[:, start:stop] + + # reshape to [num_envs, time, num layers, hidden dim] (original shape: [time, num_layers, num_envs, hidden_dim]) + # then take only time steps after dones (flattens num envs and time dimensions), + # take a batch of trajectories and finally reshape back to [num_layers, batch, hidden_dim] + last_was_done = last_was_done.permute(1, 0) + hid_a_batch = [ + saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] + .transpose(1, 0) + .contiguous() + for saved_hidden_states in self.saved_hidden_states_a + ] + hid_c_batch = [ + saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] + .transpose(1, 0) + .contiguous() + for saved_hidden_states in self.saved_hidden_states_c + ] + # remove the tuple for GRU + hid_a_batch = hid_a_batch[0] if len(hid_a_batch) == 1 else hid_a_batch + hid_c_batch = hid_c_batch[0] if len(hid_c_batch) == 1 else hid_c_batch + + yield obs_batch, actions_batch, values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( + hid_a_batch, + hid_c_batch, + ), masks_batch + + first_traj = last_traj diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/__init__.py new file mode 100644 index 0000000..f5781f1 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Helper functions.""" + +from .utils import * +from .motion_loader_g1 import G1_AMPLoader + +__all__ = [ + "G1_AMPLoader", +] \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py new file mode 100644 index 0000000..677c630 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py @@ -0,0 +1,388 @@ +import os +from os.path import join as pjoin +import glob +import json +import logging + +import torch +import numpy as np +from pybullet_utils import transformations + +from rsl_rl.utils import motion_util + +_EPS = np.finfo(float).eps * 4.0 +def quaternion_slerp(q0, q1, fraction, spin=0, shortestpath=True): + """Batch quaternion spherical linear interpolation.""" + + out = torch.zeros_like(q0) + + zero_mask = torch.isclose(fraction, torch.zeros_like(fraction)).squeeze() + ones_mask = torch.isclose(fraction, torch.ones_like(fraction)).squeeze() + out[zero_mask] = q0[zero_mask] + out[ones_mask] = q1[ones_mask] + + d = torch.sum(q0 * q1, dim=-1, keepdim=True) + dist_mask = (torch.abs(torch.abs(d) - 1.0) < _EPS).squeeze() + out[dist_mask] = q0[dist_mask] + + if shortestpath: + d_old = torch.clone(d) + d = torch.where(d_old < 0, -d, d) + q1 = torch.where(d_old < 0, -q1, q1) + + angle = torch.acos(d) + spin * torch.pi + angle_mask = (torch.abs(angle) < _EPS).squeeze() + out[angle_mask] = q0[angle_mask] + + final_mask = torch.logical_or(zero_mask, ones_mask) + final_mask = torch.logical_or(final_mask, dist_mask) + final_mask = torch.logical_or(final_mask, angle_mask) + final_mask = torch.logical_not(final_mask) + + isin = 1.0 / angle + q0 *= torch.sin((1.0 - fraction) * angle) * isin + q1 *= torch.sin(fraction * angle) * isin + q0 += q1 + out[final_mask] = q0[final_mask] + return out + + +class G1_AMPLoader: + + def __init__( + self, + device, + time_between_frames, + motion_files, + preload_transitions=False, + num_preload_transitions=1000000, + num_frames=5, + ): + """Expert dataset provides AMP observations from Dog mocap dataset. + + time_between_frames: Amount of time in seconds between transition. + """ + self.device = device + self.time_between_frames = time_between_frames + self.num_frames = num_frames + + # Values to store for each trajectory. + self.trajectories = [] + self.trajectories_full = [] + self.trajectory_names = [] + self.trajectory_idxs = [] + self.trajectory_lens = [] # Traj length in seconds. + self.trajectory_weights = [] + self.trajectory_frame_durations = [] + self.trajectory_num_frames = [] + self.motion_dir = motion_files + # import ipdb; ipdb.set_trace() + for i, motion_file in enumerate(os.listdir(motion_files)): + self.trajectory_names.append(motion_file) + motion_path = pjoin(motion_files, motion_file) + motion_data = np.load(motion_path, allow_pickle=True) + motion_data_processed = np.zeros((motion_data.shape[0],36)) + + for f_i in range(motion_data.shape[0]): + motion_data_processed[f_i, :3] = motion_data[f_i, :3] # base pos + motion_data_processed[f_i, 3:7] = motion_data[f_i, 3:7] # base quat (wxyz) + motion_data_processed[f_i, 7:35] = motion_data[f_i, 7:35] # base vel + ''' + NOTE The order of motion_data_processed is + base pos 0:3, + base quat 3:7, wxyz + dof pos 7:36, (mujoco joint order) + ''' + self.trajectories.append(torch.tensor( + motion_data_processed[:, 7:], + dtype=torch.float32, + device=self.device + )) + + self.trajectories_full.append(torch.tensor( + motion_data_processed, + dtype=torch.float32, + device=self.device + )) + + self.trajectory_idxs.append(i) + self.trajectory_weights.append(1 / len(os.listdir(motion_files))) + frame_duration = 1 / 50 + + self.trajectory_frame_durations.append(frame_duration) + traj_len = (motion_data_processed.shape[0] - 1) * frame_duration # seconds + self.trajectory_lens.append(traj_len) + self.trajectory_num_frames.append(float(motion_data_processed.shape[0])) + print(f"Loaded {traj_len}s. motion from {motion_file}.") + + # Trajectory weights are used to sample some trajectories more than others. + self.trajectory_weights = np.array(self.trajectory_weights) / np.sum(self.trajectory_weights) + self.trajectory_frame_durations = np.array(self.trajectory_frame_durations) + self.trajectory_lens = np.array(self.trajectory_lens) + self.trajectory_num_frames = np.array(self.trajectory_num_frames) + + # Preload transitions. + self.preload_transitions = preload_transitions + if self.preload_transitions: + print(f'Preloading {num_preload_transitions} transitions') + + traj_idxs = self.weighted_traj_idx_sample_batch(num_preload_transitions) + times = self.traj_time_sample_batch(traj_idxs) + self.preloaded_s_prior = self.get_full_frame_at_time_batch(traj_idxs, times - self.time_between_frames) + self.preloaded_s = self.get_full_frame_at_time_batch(traj_idxs, times) + self.preloaded_s_next = self.get_full_frame_at_time_batch(traj_idxs, times + self.time_between_frames) + print(f'Finished preloading') + + # 预加载多帧数据 + self.preloaded_frames = [] + for i in range(self.num_frames): + frame_time = times + (i - (self.num_frames - 2)) * self.time_between_frames + full_frame = self.get_full_frame_at_time_batch(traj_idxs, frame_time) + # 预处理：提前提取并连接需要的列（7:26 和 29:33），避免每次生成时重复切片 + processed_frame = torch.cat([ + full_frame[:, 7:26], + full_frame[:, 29:33] + ], dim=-1) + self.preloaded_frames.append(processed_frame) + print(f'Finished preloading multiple frames') + + self.all_trajectories_full = torch.vstack(self.trajectories_full) + + def weighted_traj_idx_sample(self): + """Get traj idx via weighted sampling.""" + return np.random.choice( + self.trajectory_idxs, p=self.trajectory_weights) + + def weighted_traj_idx_sample_batch(self, size): + """Batch sample traj idxs.""" + return np.random.choice( + self.trajectory_idxs, size=size, p=self.trajectory_weights, + replace=True) + + def traj_time_sample(self, traj_idx): + """Sample random time for traj.""" + subst = self.time_between_frames + self.trajectory_frame_durations[traj_idx] + return max( + 0, (self.trajectory_lens[traj_idx] * np.random.uniform() - subst)) + + def traj_time_sample_batch(self, traj_idxs): + """Sample random time for multiple trajectories.""" + subst = self.time_between_frames + self.trajectory_frame_durations[traj_idxs] + time_samples = self.trajectory_lens[traj_idxs] * np.random.uniform(size=len(traj_idxs)) - subst + return np.maximum(np.zeros_like(time_samples), time_samples) + + def slerp(self, val0, val1, blend): + return (1.0 - blend) * val0 + blend * val1 + + def get_trajectory(self, traj_idx): + """Returns trajectory of AMP observations.""" + return self.trajectories_full[traj_idx] + + def get_frame_at_time(self, traj_idx, time): + """Returns frame for the given trajectory at the specified time.""" + p = float(time) / self.trajectory_lens[traj_idx] + n = self.trajectories[traj_idx].shape[0] + idx_low, idx_high = int(np.floor(p * n)), int(np.ceil(p * n)) + frame_start = self.trajectories[traj_idx][idx_low] + frame_end = self.trajectories[traj_idx][idx_high] + blend = p * n - idx_low + return self.slerp(frame_start, frame_end, blend) + + def get_frame_at_time_batch(self, traj_idxs, times): + """Returns frame for the given trajectory at the specified time.""" + p = times / self.trajectory_lens[traj_idxs] + n = self.trajectory_num_frames[traj_idxs] + idx_low, idx_high = np.floor(p * n).astype(np.int32), np.ceil(p * n).astype(np.int32) + all_frame_starts = torch.zeros(len(traj_idxs), self.observation_dim, device=self.device) + all_frame_ends = torch.zeros(len(traj_idxs), self.observation_dim, device=self.device) + for traj_idx in set(traj_idxs): + trajectory = self.trajectories[traj_idx] + traj_mask = traj_idxs == traj_idx + all_frame_starts[traj_mask] = trajectory[idx_low[traj_mask]] + all_frame_ends[traj_mask] = trajectory[idx_high[traj_mask]] + blend = torch.tensor(p * n - idx_low, device=self.device, dtype=torch.float32).unsqueeze(-1) + return self.slerp(all_frame_starts, all_frame_ends, blend) + + def get_full_frame_at_time(self, traj_idx, time): + """Returns full frame for the given trajectory at the specified time.""" + p = float(time) / self.trajectory_lens[traj_idx] + n = self.trajectories_full[traj_idx].shape[0] + idx_low, idx_high = int(np.floor(p * n)), int(np.ceil(p * n)) + frame_start = self.trajectories_full[traj_idx][idx_low] + frame_end = self.trajectories_full[traj_idx][idx_high] + blend = p * n - idx_low + print(idx_low, idx_high) + return self.blend_frame_pose(frame_start, frame_end, blend) + + def get_full_frame_at_time_batch(self, traj_idxs, times): + p = times / self.trajectory_lens[traj_idxs] + n = self.trajectory_num_frames[traj_idxs] + idx_low, idx_high = np.floor(p * n).astype(np.int32), np.ceil(p * n).astype(np.int32) + all_frame_pos_starts = torch.zeros(len(traj_idxs), 3, device=self.device) + all_frame_pos_ends = torch.zeros(len(traj_idxs), 3, device=self.device) + all_frame_rot_starts = torch.zeros(len(traj_idxs), 4, device=self.device) + all_frame_rot_ends = torch.zeros(len(traj_idxs), 4, device=self.device) + all_frame_amp_starts = torch.zeros(len(traj_idxs), 29, device=self.device) + all_frame_amp_ends = torch.zeros(len(traj_idxs), 29, device=self.device) + for traj_idx in set(traj_idxs): + trajectory = self.trajectories_full[traj_idx] + traj_mask = traj_idxs == traj_idx + all_frame_pos_starts[traj_mask] = G1_AMPLoader.get_root_pos_batch(trajectory[idx_low[traj_mask]]) + all_frame_pos_ends[traj_mask] = G1_AMPLoader.get_root_pos_batch(trajectory[idx_high[traj_mask]]) + all_frame_rot_starts[traj_mask] = G1_AMPLoader.get_root_rot_batch(trajectory[idx_low[traj_mask]]) + all_frame_rot_ends[traj_mask] = G1_AMPLoader.get_root_rot_batch(trajectory[idx_high[traj_mask]]) + all_frame_amp_starts[traj_mask] = trajectory[idx_low[traj_mask]][:, 7:36] # base vel3+ang3, dof vel23+ang23 + all_frame_amp_ends[traj_mask] = trajectory[idx_high[traj_mask]][:, 7:36] # base vel3+ang3, dof vel23+ang23 + blend = torch.tensor(p * n - idx_low, device=self.device, dtype=torch.float32).unsqueeze(-1) + pos_blend = self.slerp(all_frame_pos_starts, all_frame_pos_ends, blend) + rot_blend = quaternion_slerp(all_frame_rot_starts, all_frame_rot_ends, blend) + amp_blend = self.slerp(all_frame_amp_starts, all_frame_amp_ends, blend) + return torch.cat([pos_blend, rot_blend, amp_blend], dim=-1) + + def get_frame(self): + """Returns random frame.""" + traj_idx = self.weighted_traj_idx_sample() + sampled_time = self.traj_time_sample(traj_idx) + return self.get_frame_at_time(traj_idx, sampled_time) + + def get_full_frame(self): + """Returns random full frame.""" + traj_idx = self.weighted_traj_idx_sample() + sampled_time = self.traj_time_sample(traj_idx) + return self.get_full_frame_at_time(traj_idx, sampled_time) + + def get_full_frame_batch(self, num_frames): + if self.preload_transitions: + idxs = np.random.choice( + self.preloaded_s.shape[0], size=num_frames) + return self.preloaded_s[idxs] + else: + traj_idxs = self.weighted_traj_idx_sample_batch(num_frames) + times = self.traj_time_sample_batch(traj_idxs) + return self.get_full_frame_at_time_batch(traj_idxs, times) + + def blend_frame_pose(self, frame0, frame1, blend): + """Linearly interpolate between two frames, including orientation. + + Args: + frame0: First frame to be blended corresponds to (blend = 0). + frame1: Second frame to be blended corresponds to (blend = 1). + blend: Float between [0, 1], specifying the interpolation between + the two frames. + Returns: + An interpolation of the two frames. + """ + root_pos0, root_pos1 = G1_AMPLoader.get_root_pos(frame0), G1_AMPLoader.get_root_pos(frame1) + root_rot0, root_rot1 = G1_AMPLoader.get_root_rot(frame0), G1_AMPLoader.get_root_rot(frame1) + joints0, joints1 = G1_AMPLoader.get_joint_pose(frame0), G1_AMPLoader.get_joint_pose(frame1) + # tar_toe_pos_0, tar_toe_pos_1 = G1_AMPLoader.get_tar_toe_pos_local(frame0), G1_AMPLoader.get_tar_toe_pos_local(frame1) + linear_vel_0, linear_vel_1 = G1_AMPLoader.get_linear_vel(frame0), G1_AMPLoader.get_linear_vel(frame1) + angular_vel_0, angular_vel_1 = G1_AMPLoader.get_angular_vel(frame0), G1_AMPLoader.get_angular_vel(frame1) + joint_vel_0, joint_vel_1 = G1_AMPLoader.get_joint_vel(frame0), G1_AMPLoader.get_joint_vel(frame1) + + blend_root_pos = self.slerp(root_pos0, root_pos1, blend) + blend_root_rot = transformations.quaternion_slerp(root_rot0.cpu().numpy(), root_rot1.cpu().numpy(), blend) + blend_root_rot = torch.tensor(motion_util.standardize_quaternion(blend_root_rot),dtype=torch.float32, device=self.device) + blend_joints = self.slerp(joints0, joints1, blend) + # blend_tar_toe_pos = self.slerp(tar_toe_pos_0, tar_toe_pos_1, blend) + blend_linear_vel = self.slerp(linear_vel_0, linear_vel_1, blend) + blend_angular_vel = self.slerp(angular_vel_0, angular_vel_1, blend) + blend_joints_vel = self.slerp(joint_vel_0, joint_vel_1, blend) + + # return + # torch.cat([ + # blend_root_pos, blend_root_rot, blend_linear_vel, blend_angular_vel, blend_joints, blend_joints_vel]) + return torch.cat([blend_root_pos, blend_root_rot, blend_linear_vel, blend_angular_vel, blend_joints]) + + def feed_forward_generator_23dof_multi(self, num_mini_batch, mini_batch_size): + """Generates a batch of AMP transitions.""" + # import ipdb; ipdb.set_trace() + for _ in range(num_mini_batch): + if self.preload_transitions: + idxs = np.random.choice(self.preloaded_s.shape[0], size=mini_batch_size) + + frames = [] + for i in range(self.num_frames): + # 数据已在预加载时预处理，直接索引即可 + s = self.preloaded_frames[i][idxs] + frames.append(s) + else: + NotImplementedError('preload transition') + yield torch.stack(frames, dim=1) # [batch, num_frames, 16] + + + + + def quaternion_to_euler_array(self, quat): + # Ensure quaternion is in the correct format [x, y, z, w] + x, y, z, w =quat + + # Roll (x-axis rotation) + t0 = +2.0 * (w * x + y * z) + t1 = +1.0 - 2.0 * (x * x + y * y) + roll_x = np.arctan2(t0, t1) + + # Pitch (y-axis rotation) + t2 = +2.0 * (w * y - z * x) + t2 = np.clip(t2, -1.0, 1.0) + pitch_y = np.arcsin(t2) + + # Yaw (z-axis rotation) + t3 = +2.0 * (w * z + x * y) + t4 = +1.0 - 2.0 * (y * y + z * z) + yaw_z = np.arctan2(t3, t4) + + # Returns roll, pitch, yaw in a NumPy array in radians + return np.array([roll_x, pitch_y, yaw_z]) + + def euler_to_quaternion(self, root_rot): + roll, pitch, yaw = root_rot[0], root_rot[1], root_rot[2] + cy = np.cos(yaw * 0.5) + sy = np.sin(yaw * 0.5) + cp = np.cos(pitch * 0.5) + sp = np.sin(pitch * 0.5) + cr = np.cos(roll * 0.5) + sr = np.sin(roll * 0.5) + + qw = cy * cp * cr + sy * sp * sr + qx = cy * cp * sr - sy * sp * cr + qy = sy * cp * sr + cy * sp * cr + qz = sy * cp * cr - cy * sp * sr + + return np.array([qx, qy, qz, qw]) + + @property + def observation_dim(self): + """Size of AMP observations.""" + return self.trajectories[0].shape[1] + 1 + + @property + def num_motions(self): + return len(self.trajectory_names) + @staticmethod + def get_root_pos(pose): + return pose[0:3] + + @staticmethod + def get_root_pos_batch(poses): + return poses[:, 0:3] + + @staticmethod + def get_root_rot(pose): + return pose[3:7] + + @staticmethod + def get_root_rot_batch(poses): + return poses[:, 3:7] + + @staticmethod + def get_joint_pose_batch_12dof(poses): + return poses[:, 13:25] + + @staticmethod + def get_tar_toe_pos_local(pose): + return pose[G1_AMPLoader.TAR_TOE_POS_LOCAL_START_IDX:G1_AMPLoader.TAR_TOE_POS_LOCAL_END_IDX] + + @staticmethod + def get_tar_toe_pos_local_batch(poses): + return poses[:, G1_AMPLoader.TAR_TOE_POS_LOCAL_START_IDX:G1_AMPLoader.TAR_TOE_POS_LOCAL_END_IDX] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_util.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_util.py new file mode 100644 index 0000000..3d49bfa --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_util.py @@ -0,0 +1,97 @@ +# coding=utf-8 +# Copyright 2020 The Google Research Authors. +# +# 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. + +"""Utility functions for processing motion clips.""" + +import os +import inspect +currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) +parentdir = os.path.dirname(os.path.dirname(currentdir)) +os.sys.path.insert(0, parentdir) + +import numpy as np + +from rsl_rl.utils import pose3d +# from pybullet_utils import transformations + + +def standardize_quaternion(q): + """Returns a quaternion where q.w >= 0 to remove redundancy due to q = -q. + + Args: + q: A quaternion to be standardized. + + Returns: + A quaternion with q.w >= 0. + + """ + if q[-1] < 0: + q = -q + return q + + +def normalize_rotation_angle(theta): + """Returns a rotation angle normalized between [-pi, pi]. + + Args: + theta: angle of rotation (radians). + + Returns: + An angle of rotation normalized between [-pi, pi]. + + """ + norm_theta = theta + if np.abs(norm_theta) > np.pi: + norm_theta = np.fmod(norm_theta, 2 * np.pi) + if norm_theta >= 0: + norm_theta += -2 * np.pi + else: + norm_theta += 2 * np.pi + + return norm_theta + + +def calc_heading(q): + """Returns the heading of a rotation q, specified as a quaternion. + + The heading represents the rotational component of q along the vertical + axis (z axis). + + Args: + q: A quaternion that the heading is to be computed from. + + Returns: + An angle representing the rotation about the z axis. + + """ + ref_dir = np.array([1, 0, 0]) + rot_dir = pose3d.QuaternionRotatePoint(ref_dir, q) + heading = np.arctan2(rot_dir[1], rot_dir[0]) + return heading + + +# def calc_heading_rot(q): +# """Return a quaternion representing the heading rotation of q along the vertical axis (z axis). + +# Args: +# q: A quaternion that the heading is to be computed from. + +# Returns: +# A quaternion representing the rotation about the z axis. + +# """ +# heading = calc_heading(q) +# q_heading = transformations.quaternion_about_axis(heading, [0, 0, 1]) +# return q_heading diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py new file mode 100644 index 0000000..3796ec8 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py @@ -0,0 +1,94 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +from dataclasses import asdict +from torch.utils.tensorboard import SummaryWriter + +try: + import neptune +except ModuleNotFoundError: + raise ModuleNotFoundError("neptune-client is required to log to Neptune.") + + +class NeptuneLogger: + def __init__(self, project, token): + self.run = neptune.init_run(project=project, api_token=token) + + def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): + self.run["runner_cfg"] = runner_cfg + self.run["policy_cfg"] = policy_cfg + self.run["alg_cfg"] = alg_cfg + self.run["env_cfg"] = asdict(env_cfg) + + +class NeptuneSummaryWriter(SummaryWriter): + """Summary writer for Neptune.""" + + def __init__(self, log_dir: str, flush_secs: int, cfg): + super().__init__(log_dir, flush_secs) + + try: + project = cfg["neptune_project"] + except KeyError: + raise KeyError("Please specify neptune_project in the runner config, e.g. legged_gym.") + + try: + token = os.environ["NEPTUNE_API_TOKEN"] + except KeyError: + raise KeyError( + "Neptune api token not found. Please run or add to ~/.bashrc: export NEPTUNE_API_TOKEN=YOUR_API_TOKEN" + ) + + try: + entity = os.environ["NEPTUNE_USERNAME"] + except KeyError: + raise KeyError( + "Neptune username not found. Please run or add to ~/.bashrc: export NEPTUNE_USERNAME=YOUR_USERNAME" + ) + + neptune_project = entity + "/" + project + + self.neptune_logger = NeptuneLogger(neptune_project, token) + + self.name_map = { + "Train/mean_reward/time": "Train/mean_reward_time", + "Train/mean_episode_length/time": "Train/mean_episode_length_time", + } + + run_name = os.path.split(log_dir)[-1] + + self.neptune_logger.run["log_dir"].log(run_name) + + def _map_path(self, path): + if path in self.name_map: + return self.name_map[path] + else: + return path + + def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_style=False): + super().add_scalar( + tag, + scalar_value, + global_step=global_step, + walltime=walltime, + new_style=new_style, + ) + self.neptune_logger.run[self._map_path(tag)].log(scalar_value, step=global_step) + + def stop(self): + self.neptune_logger.run.stop() + + def log_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): + self.neptune_logger.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg) + + def save_model(self, model_path, iter): + self.neptune_logger.run["model/saved_model_" + str(iter)].upload(model_path) + + def save_file(self, path, iter=None): + name = path.rsplit("/", 1)[-1].split(".")[0] + self.neptune_logger.run["git_diff/" + name].upload(path) diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/pose3d.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/pose3d.py new file mode 100644 index 0000000..a4cba1e --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/pose3d.py @@ -0,0 +1,283 @@ +# coding=utf-8 +# Copyright 2020 The Google Research Authors. +# +# 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. +"""Utilities for 3D pose conversion.""" +import math +import numpy as np + +# from pybullet_utils import transformations + +VECTOR3_0 = np.zeros(3, dtype=np.float64) +VECTOR3_1 = np.ones(3, dtype=np.float64) +VECTOR3_X = np.array([1, 0, 0], dtype=np.float64) +VECTOR3_Y = np.array([0, 1, 0], dtype=np.float64) +VECTOR3_Z = np.array([0, 0, 1], dtype=np.float64) + +# QUATERNION_IDENTITY is the multiplicative identity 1.0 + 0i + 0j + 0k. +# When interpreted as a rotation, it is the identity rotation. +QUATERNION_IDENTITY = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float64) + + +def Vector3RandomNormal(sigma, mu=VECTOR3_0): + """Returns a random 3D vector from a normal distribution. + + Each component is selected independently from a normal distribution. + + Args: + sigma: Scale (or stddev) of distribution for all variables. + mu: Mean of distribution for each variable. + + Returns: + A 3D vector in a numpy array. + """ + + random_v3 = np.random.normal(scale=sigma, size=3) + mu + return random_v3 + + +def Vector3RandomUniform(low=VECTOR3_0, high=VECTOR3_1): + """Returns a 3D vector selected uniformly from the input box. + + Args: + low: The min-value corner of the box. + high: The max-value corner of the box. + + Returns: + A 3D vector in a numpy array. + """ + + random_x = np.random.uniform(low=low[0], high=high[0]) + random_y = np.random.uniform(low=low[1], high=high[1]) + random_z = np.random.uniform(low=low[2], high=high[2]) + return np.array([random_x, random_y, random_z]) + + +def Vector3RandomUnit(): + """Returns a random 3D vector with unit length. + + Generates a 3D vector selected uniformly from the unit sphere. + + Returns: + A normalized 3D vector in a numpy array. + """ + longitude = np.random.uniform(low=-math.pi, high=math.pi) + sin_latitude = np.random.uniform(low=-1.0, high=1.0) + cos_latitude = math.sqrt(1.0 - sin_latitude * sin_latitude) + x = math.cos(longitude) * cos_latitude + y = math.sin(longitude) * cos_latitude + z = sin_latitude + return np.array([x, y, z], dtype=np.float64) + + +def QuaternionNormalize(q): + """Normalizes the quaternion to length 1. + + Divides the quaternion by its magnitude. If the magnitude is too + small, returns the quaternion identity value (1.0). + + Args: + q: A quaternion to be normalized. + + Raises: + ValueError: If input quaternion has length near zero. + + Returns: + A quaternion with magnitude 1 in a numpy array [x, y, z, w]. + + """ + q_norm = np.linalg.norm(q) + if np.isclose(q_norm, 0.0): + raise ValueError( + 'Quaternion may not be zero in QuaternionNormalize: |q| = %f, q = %s' % + (q_norm, q)) + return q / q_norm + + +def QuaternionFromAxisAngle(axis, angle): + """Returns a quaternion that generates the given axis-angle rotation. + + Returns the quaternion: sin(angle/2) * axis + cos(angle/2). + + Args: + axis: Axis of rotation, a 3D vector in a numpy array. + angle: The angle of rotation (radians). + + Raises: + ValueError: If input axis is not a normalizable 3D vector. + + Returns: + A unit quaternion in a numpy array. + + """ + if len(axis) != 3: + raise ValueError('Axis vector should have three components: %s' % axis) + axis_norm = np.linalg.norm(axis) + if np.isclose(axis_norm, 0.0): + raise ValueError('Axis vector may not have zero length: |v| = %f, v = %s' % + (axis_norm, axis)) + half_angle = angle * 0.5 + q = np.zeros(4, dtype=np.float64) + q[0:3] = axis + q[0:3] *= math.sin(half_angle) / axis_norm + q[3] = math.cos(half_angle) + return q + + +def QuaternionToAxisAngle(quat, default_axis=VECTOR3_Z, direction_axis=None): + """Calculates axis and angle of rotation performed by a quaternion. + + Calculates the axis and angle of the rotation performed by the quaternion. + The quaternion should have four values and be normalized. + + Args: + quat: Unit quaternion in a numpy array. + default_axis: 3D vector axis used if the rotation is near to zero. Without + this default, small rotations would result in an exception. It is + reasonable to use a default axis for tiny rotations, because zero angle + rotations about any axis are equivalent. + direction_axis: Used to disambiguate rotation directions. If the + direction_axis is specified, the axis of the rotation will be chosen such + that its inner product with the direction_axis is non-negative. + + Raises: + ValueError: If quat is not a normalized quaternion. + + Returns: + axis: Axis of rotation. + angle: Angle in radians. + """ + if len(quat) != 4: + raise ValueError( + 'Quaternion should have four components [x, y, z, w]: %s' % quat) + if not np.isclose(1.0, np.linalg.norm(quat)): + raise ValueError('Quaternion should have unit length: |q| = %f, q = %s' % + (np.linalg.norm(quat), quat)) + axis = quat[:3].copy() + axis_norm = np.linalg.norm(axis) + min_axis_norm = 1e-8 + if axis_norm < min_axis_norm: + axis = default_axis + if len(default_axis) != 3: + raise ValueError('Axis vector should have three components: %s' % axis) + if not np.isclose(np.linalg.norm(axis), 1.0): + raise ValueError('Axis vector should have unit length: |v| = %f, v = %s' % + (np.linalg.norm(axis), axis)) + else: + axis /= axis_norm + sin_half_angle = axis_norm + if direction_axis is not None and np.inner(axis, direction_axis) < 0: + sin_half_angle = -sin_half_angle + axis = -axis + cos_half_angle = quat[3] + half_angle = math.atan2(sin_half_angle, cos_half_angle) + angle = half_angle * 2 + return axis, angle + + +def QuaternionRandomRotation(max_angle=math.pi): + """Creates a random small rotation around a random axis. + + Generates a small rotation with the axis vector selected uniformly + from the unit sphere and an angle selected from a uniform + distribution over [0, max_angle]. + + If the max_angle is not specified, the rotation should be selected + uniformly over all possible rotation angles. + + Args: + max_angle: The maximum angle of rotation (radians). + + Returns: + A unit quaternion in a numpy array. + + """ + + angle = np.random.uniform(low=0, high=max_angle) + axis = Vector3RandomUnit() + return QuaternionFromAxisAngle(axis, angle) + + +# def QuaternionRotatePoint(point, quat): +# """Performs a rotation by quaternion. + +# Rotate the point by the quaternion using quaternion multiplication, +# (q * p * q^-1), without constructing the rotation matrix. + +# Args: +# point: The point to be rotated. +# quat: The rotation represented as a quaternion [x, y, z, w]. + +# Returns: +# A 3D vector in a numpy array. +# """ + +# q_point = np.array([point[0], point[1], point[2], 0.0]) +# quat_inverse = transformations.quaternion_inverse(quat) +# q_point_rotated = transformations.quaternion_multiply( +# transformations.quaternion_multiply(quat, q_point), quat_inverse) +# return q_point_rotated[:3] + + +def IsRotationMatrix(m): + """Returns true if the 3x3 submatrix represents a rotation. + + Args: + m: A transformation matrix. + + Raises: + ValueError: If input is not a matrix of size at least 3x3. + + Returns: + True if the 3x3 submatrix is a rotation (orthogonal). + """ + if len(m.shape) != 2 or m.shape[0] < 3 or m.shape[1] < 3: + raise ValueError('Matrix should be 3x3 or 4x4: %s\n %s' % (m.shape, m)) + rot = m[:3, :3] + eye = np.matmul(rot, np.transpose(rot)) + return np.isclose(eye, np.identity(3), atol=1e-4).all() + +# def ZAxisAlignedRobotPoseTool(robot_pose_tool): +# """Returns the current gripper pose rotated for alignment with the z-axis. + +# Args: +# robot_pose_tool: a pose3d.Pose3d() instance. + +# Returns: +# An instance of pose.Transform representing the current gripper pose +# rotated for alignment with the z-axis. +# """ +# # Align the current pose to the z-axis. +# robot_pose_tool.quaternion = transformations.quaternion_multiply( +# RotationBetween( +# robot_pose_tool.matrix4x4[0:3, 0:3].dot(np.array([0, 0, 1])), +# np.array([0.0, 0.0, -1.0])), robot_pose_tool.quaternion) +# return robot_pose_tool + +# def RotationBetween(a_translation_b, a_translation_c): +# """Computes the rotation from one vector to another. + +# The computed rotation has the property that: + +# a_translation_c = a_rotation_b_to_c * a_translation_b + +# Args: +# a_translation_b: vec3, vector to rotate from +# a_translation_c: vec3, vector to rotate to + +# Returns: +# a_rotation_b_to_c: new Orientation +# """ +# rotation = rotation3.Rotation3.rotation_between( +# a_translation_b, a_translation_c, err_msg='RotationBetween') +# return rotation.quaternion.xyzw diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/utils.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/utils.py new file mode 100644 index 0000000..3605622 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/utils.py @@ -0,0 +1,360 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import git +import importlib +import os +import pathlib +import torch +import warnings +from tensordict import TensorDict +from typing import Callable +import numpy as np +class RunningMeanStd: + def __init__(self, epsilon: float = 1e-4, shape: Tuple[int, ...] = ()): + """ + Calculates the running mean and std of a data stream + https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm + :param epsilon: helps with arithmetic issues + :param shape: the shape of the data stream's output + """ + self.mean = np.zeros(shape, np.float64) + self.var = np.ones(shape, np.float64) + self.count = epsilon + + def update(self, arr: np.ndarray) -> None: + batch_mean = np.mean(arr, axis=0) + batch_var = np.var(arr, axis=0) + batch_count = arr.shape[0] + self.update_from_moments(batch_mean, batch_var, batch_count) + + def update_from_moments(self, batch_mean: np.ndarray, batch_var: np.ndarray, batch_count: int) -> None: + delta = batch_mean - self.mean + tot_count = self.count + batch_count + + new_mean = self.mean + delta * batch_count / tot_count + m_a = self.var * self.count + m_b = batch_var * batch_count + m_2 = m_a + m_b + np.square(delta) * self.count * batch_count / (self.count + batch_count) + new_var = m_2 / (self.count + batch_count) + + new_count = batch_count + self.count + + self.mean = new_mean + self.var = new_var + self.count = new_count + + +class Normalizer(RunningMeanStd): + def __init__(self, input_dim, epsilon=1e-4, clip_obs=10.0): + super().__init__(shape=input_dim) + self.epsilon = epsilon + self.clip_obs = clip_obs + + def normalize(self, input): + return np.clip((input - self.mean) / np.sqrt(self.var + self.epsilon), -self.clip_obs, self.clip_obs) + + def normalize_torch(self, input, device): + mean_torch = torch.tensor(self.mean, device=device, dtype=torch.float32) + std_torch = torch.sqrt(torch.tensor(self.var + self.epsilon, device=device, dtype=torch.float32)) + return torch.clamp((input - mean_torch) / std_torch, -self.clip_obs, self.clip_obs) + + def update_normalizer(self, rollouts, expert_loader): + policy_data_generator = rollouts.feed_forward_generator_amp(None, mini_batch_size=expert_loader.batch_size) + expert_data_generator = expert_loader.dataset.feed_forward_generator_amp(expert_loader.batch_size) + + for expert_batch, policy_batch in zip(expert_data_generator, policy_data_generator): + self.update(torch.vstack(tuple(policy_batch) + tuple(expert_batch)).cpu().numpy()) + + +def resolve_nn_activation(act_name: str) -> torch.nn.Module: + """Resolves the activation function from the name. + + Args: + act_name: The name of the activation function. + + Returns: + The activation function. + + Raises: + ValueError: If the activation function is not found. + """ + act_dict = { + "elu": torch.nn.ELU(), + "selu": torch.nn.SELU(), + "relu": torch.nn.ReLU(), + "crelu": torch.nn.CELU(), + "lrelu": torch.nn.LeakyReLU(), + "tanh": torch.nn.Tanh(), + "sigmoid": torch.nn.Sigmoid(), + "softplus": torch.nn.Softplus(), + "gelu": torch.nn.GELU(), + "swish": torch.nn.SiLU(), + "mish": torch.nn.Mish(), + "identity": torch.nn.Identity(), + } + + act_name = act_name.lower() + if act_name in act_dict: + return act_dict[act_name] + else: + raise ValueError(f"Invalid activation function '{act_name}'. Valid activations are: {list(act_dict.keys())}") + + +def resolve_optimizer(optimizer_name: str) -> torch.optim.Optimizer: + """Resolves the optimizer from the name. + + Args: + optimizer_name: The name of the optimizer. + + Returns: + The optimizer. + + Raises: + ValueError: If the optimizer is not found. + """ + optimizer_dict = { + "adam": torch.optim.Adam, + "adamw": torch.optim.AdamW, + "sgd": torch.optim.SGD, + "rmsprop": torch.optim.RMSprop, + } + + optimizer_name = optimizer_name.lower() + if optimizer_name in optimizer_dict: + return optimizer_dict[optimizer_name] + else: + raise ValueError(f"Invalid optimizer '{optimizer_name}'. Valid optimizers are: {list(optimizer_dict.keys())}") + + +def split_and_pad_trajectories( + tensor: torch.Tensor | TensorDict, dones: torch.Tensor +) -> tuple[torch.Tensor | TensorDict, torch.Tensor]: + """Splits trajectories at done indices. Then concatenates them and pads with zeros up to the length of the longest + trajectory. Returns masks corresponding to valid parts of the trajectories. + + Example: + Input: [[a1, a2, a3, a4 | a5, a6], + [b1, b2 | b3, b4, b5 | b6]] + + Output:[[a1, a2, a3, a4], | [[True, True, True, True], + [a5, a6, 0, 0], | [True, True, False, False], + [b1, b2, 0, 0], | [True, True, False, False], + [b3, b4, b5, 0], | [True, True, True, False], + [b6, 0, 0, 0]] | [True, False, False, False]] + + Assumes that the input has the following order of dimensions: [time, number of envs, additional dimensions] + """ + + dones = dones.clone() + dones[-1] = 1 + # Permute the buffers to have order (num_envs, num_transitions_per_env, ...), for correct reshaping + flat_dones = dones.transpose(1, 0).reshape(-1, 1) + # Get length of trajectory by counting the number of successive not done elements + done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero()[:, 0])) + trajectory_lengths = done_indices[1:] - done_indices[:-1] + trajectory_lengths_list = trajectory_lengths.tolist() + # Extract the individual trajectories + if isinstance(tensor, TensorDict): + padded_trajectories = {} + for k, v in tensor.items(): + # split the tensor into trajectories + trajectories = torch.split(v.transpose(1, 0).flatten(0, 1), trajectory_lengths_list) + # add at least one full length trajectory + trajectories = trajectories + (torch.zeros(v.shape[0], *v.shape[2:], device=v.device),) + # pad the trajectories to the length of the longest trajectory + padded_trajectories[k] = torch.nn.utils.rnn.pad_sequence(trajectories) + # remove the added tensor + padded_trajectories[k] = padded_trajectories[k][:, :-1] + padded_trajectories = TensorDict( + padded_trajectories, batch_size=[tensor.batch_size[0], len(trajectory_lengths_list)] + ) + else: + # split the tensor into trajectories + trajectories = torch.split(tensor.transpose(1, 0).flatten(0, 1), trajectory_lengths_list) + # add at least one full length trajectory + trajectories = trajectories + (torch.zeros(tensor.shape[0], *tensor.shape[2:], device=tensor.device),) + # pad the trajectories to the length of the longest trajectory + padded_trajectories = torch.nn.utils.rnn.pad_sequence(trajectories) + # remove the added tensor + padded_trajectories = padded_trajectories[:, :-1] + # create masks for the valid parts of the trajectories + trajectory_masks = trajectory_lengths > torch.arange(0, tensor.shape[0], device=tensor.device).unsqueeze(1) + return padded_trajectories, trajectory_masks + + +def unpad_trajectories(trajectories, masks): + """Does the inverse operation of split_and_pad_trajectories()""" + # Need to transpose before and after the masking to have proper reshaping + return ( + trajectories.transpose(1, 0)[masks.transpose(1, 0)] + .view(-1, trajectories.shape[0], trajectories.shape[-1]) + .transpose(1, 0) + ) + + +def store_code_state(logdir, repositories) -> list: + git_log_dir = os.path.join(logdir, "git") + os.makedirs(git_log_dir, exist_ok=True) + file_paths = [] + for repository_file_path in repositories: + try: + repo = git.Repo(repository_file_path, search_parent_directories=True) + t = repo.head.commit.tree + except Exception: + print(f"Could not find git repository in {repository_file_path}. Skipping.") + # skip if not a git repository + continue + # get the name of the repository + repo_name = pathlib.Path(repo.working_dir).name + diff_file_name = os.path.join(git_log_dir, f"{repo_name}.diff") + # check if the diff file already exists + if os.path.isfile(diff_file_name): + continue + # write the diff file + print(f"Storing git diff for '{repo_name}' in: {diff_file_name}") + with open(diff_file_name, "x", encoding="utf-8") as f: + content = f"--- git status ---\n{repo.git.status()} \n\n\n--- git diff ---\n{repo.git.diff(t)}" + f.write(content) + # add the file path to the list of files to be uploaded + file_paths.append(diff_file_name) + return file_paths + + +def string_to_callable(name: str) -> Callable: + """Resolves the module and function names to return the function. + + Args: + name: The function name. The format should be 'module:attribute_name'. + + Raises: + ValueError: When the resolved attribute is not a function. + ValueError: When unable to resolve the attribute. + + Returns: + The function loaded from the module. + """ + try: + mod_name, attr_name = name.split(":") + mod = importlib.import_module(mod_name) + callable_object = getattr(mod, attr_name) + # check if attribute is callable + if callable(callable_object): + return callable_object + else: + raise ValueError(f"The imported object is not callable: '{name}'") + except AttributeError as e: + msg = ( + "We could not interpret the entry as a callable object. The format of input should be" + f" 'module:attribute_name'\nWhile processing input '{name}', received the error:\n {e}." + ) + raise ValueError(msg) + + +def resolve_obs_groups( + obs: TensorDict, obs_groups: dict[str, list[str]], default_sets: list[str] +) -> dict[str, list[str]]: + """Validates the observation configuration and defaults missing observation sets. + + The input is an observation dictionary `obs` containing observation groups and a configuration dictionary + `obs_groups` where the keys are the observation sets and the values are lists of observation groups. + + The configuration dictionary could for example look like: + { + "policy": ["group_1", "group_2"], + "critic": ["group_1", "group_3"] + } + + This means that the 'policy' observation set will contain the observations "group_1" and "group_2" and the + 'critic' observation set will contain the observations "group_1" and "group_3". This function will check that all + the observations in the 'policy' and 'critic' observation sets are present in the observation dictionary from the + environment. + + Additionally, if one of the `default_sets`, e.g. "critic", is not present in the configuration dictionary, + this function will: + + 1. Check if a group with the same name exists in the observations and assign this group to the observation set. + 2. If 1. fails, it will assign the observations from the 'policy' observation set to the default observation set. + + Args: + obs: Observations from the environment in the form of a dictionary. + obs_groups: Observation sets configuration. + default_sets: Reserved observation set names used by the algorithm (besides 'policy'). + If not provided in 'obs_groups', a default behavior gets triggered. + + Returns: + The resolved observation groups. + + Raises: + ValueError: If any observation set is an empty list. + ValueError: If any observation set contains an observation term that is not present in the observations. + """ + # check if policy observation set exists + if "policy" not in obs_groups.keys(): + if "policy" in obs: + obs_groups["policy"] = ["policy"] + warnings.warn( + "The observation configuration dictionary 'obs_groups' must contain the 'policy' key." + " As an observation group with the name 'policy' was found, this is assumed to be the observation set." + " Consider adding the 'policy' key to the 'obs_groups' dictionary for clarity." + " This behavior will be removed in a future version." + ) + else: + raise ValueError( + "The observation configuration dictionary 'obs_groups' must contain the 'policy' key." + f" Found keys: {list(obs_groups.keys())}" + ) + + # check all observation sets for valid observation groups + for set_name, groups in obs_groups.items(): + # check if the list is empty + if len(groups) == 0: + msg = f"The '{set_name}' key in the 'obs_groups' dictionary can not be an empty list." + if set_name in default_sets: + if set_name not in obs: + msg += " Consider removing the key to default to the observations used for the 'policy' set." + else: + msg += ( + f" Consider removing the key to default to the observation '{set_name}' from the environment." + ) + raise ValueError(msg) + # check groups exist inside the observations from the environment + for group in groups: + if group not in obs: + raise ValueError( + f"Observation '{group}' in observation set '{set_name}' not found in the observations from the" + f" environment. Available observations from the environment: {list(obs.keys())}" + ) + + # fill missing observation sets + for default_set_name in default_sets: + if default_set_name not in obs_groups.keys(): + if default_set_name in obs: + obs_groups[default_set_name] = [default_set_name] + warnings.warn( + f"The observation configuration dictionary 'obs_groups' must contain the '{default_set_name}' key." + f" As an observation group with the name '{default_set_name}' was found, this is assumed to be the" + f" observation set. Consider adding the '{default_set_name}' key to the 'obs_groups' dictionary for" + " clarity. This behavior will be removed in a future version." + ) + else: + obs_groups[default_set_name] = obs_groups["policy"].copy() + warnings.warn( + f"The observation configuration dictionary 'obs_groups' must contain the '{default_set_name}' key." + f" As the configuration for '{default_set_name}' is missing, the observations from the 'policy' set" + f" are used. Consider adding the '{default_set_name}' key to the 'obs_groups' dictionary for" + " clarity. This behavior will be removed in a future version." + ) + + # print the final parsed observation sets + print("-" * 80) + print("Resolved observation sets: ") + for set_name, groups in obs_groups.items(): + print("\t", set_name, ": ", groups) + print("-" * 80) + + return obs_groups diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py new file mode 100644 index 0000000..243e82d --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py @@ -0,0 +1,87 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +from dataclasses import asdict +from torch.utils.tensorboard import SummaryWriter + +try: + import wandb +except ModuleNotFoundError: + raise ModuleNotFoundError("Wandb is required to log to Weights and Biases.") + + +class WandbSummaryWriter(SummaryWriter): + """Summary writer for Weights and Biases.""" + + def __init__(self, log_dir: str, flush_secs: int, cfg): + super().__init__(log_dir, flush_secs) + + # Get the run name + run_name = os.path.split(log_dir)[-1] + + try: + project = cfg["wandb_project"] + except KeyError: + raise KeyError("Please specify wandb_project in the runner config, e.g. legged_gym.") + + try: + entity = os.environ["WANDB_USERNAME"] + except KeyError: + entity = None + + # Initialize wandb + wandb.init(project=project, entity=entity, name=run_name) + + # Add log directory to wandb + wandb.config.update({"log_dir": log_dir}) + + self.name_map = { + "Train/mean_reward/time": "Train/mean_reward_time", + "Train/mean_episode_length/time": "Train/mean_episode_length_time", + } + + def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): + wandb.config.update({"runner_cfg": runner_cfg}) + wandb.config.update({"policy_cfg": policy_cfg}) + wandb.config.update({"alg_cfg": alg_cfg}) + try: + wandb.config.update({"env_cfg": env_cfg.to_dict()}) + except Exception: + wandb.config.update({"env_cfg": asdict(env_cfg)}) + + def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_style=False): + super().add_scalar( + tag, + scalar_value, + global_step=global_step, + walltime=walltime, + new_style=new_style, + ) + wandb.log({self._map_path(tag): scalar_value}, step=global_step) + + def stop(self): + wandb.finish() + + def log_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): + self.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg) + + def save_model(self, model_path, iter): + wandb.save(model_path, base_path=os.path.dirname(model_path)) + + def save_file(self, path, iter=None): + wandb.save(path, base_path=os.path.dirname(path)) + + """ + Private methods. + """ + + def _map_path(self, path): + if path in self.name_map: + return self.name_map[path] + else: + return path diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/__init__.py new file mode 100644 index 0000000..ab7c056 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Submodule defining the environment definitions.""" + +from .vec_env import VecEnv + +__all__ = ["VecEnv"] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/vec_env.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/vec_env.py new file mode 100644 index 0000000..ae471cf --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/vec_env.py @@ -0,0 +1,113 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from abc import ABC, abstractmethod +from tensordict import TensorDict + + +class VecEnv(ABC): + """Abstract class for a vectorized environment. + + The vectorized environment is a collection of environments that are synchronized. This means that + the same type of action is applied to all environments and the same type of observation is returned from all + environments. + """ + + num_envs: int + """Number of environments.""" + + num_actions: int + """Number of actions.""" + + max_episode_length: int | torch.Tensor + + max_episode_length_s: float + """Maximum episode length. + + The maximum episode length can be a scalar or a tensor. If it is a scalar, it is the same for all environments. + If it is a tensor, it is the maximum episode length for each environment. This is useful for dynamic episode + lengths. + """ + + episode_length_buf: torch.Tensor + """Buffer for current episode lengths.""" + + device: torch.device | str + """Device to use.""" + + cfg: dict | object + """Configuration object.""" + + reset_env_ids: torch.Tensor | None = None + + contact_phase: torch.Tensor | None = None + """ + Operations. + """ + + @abstractmethod + def get_observations(self) -> TensorDict: + """Return the current observations. + + Returns: + observations (TensorDict): Observations from the environment. + """ + raise NotImplementedError + + @abstractmethod + def get_amp_observations(self) -> TensorDict: + """Return the current AMP observations. + + Returns: + observations (TensorDict): Observations from the environment. + """ + raise NotImplementedError + + @abstractmethod + def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]: + """Apply input action to the environment. + + Args: + actions (torch.Tensor): Input actions to apply. Shape: (num_envs, num_actions) + + Returns: + observations (TensorDict): Observations from the environment. + rewards (torch.Tensor): Rewards from the environment. Shape: (num_envs,) + dones (torch.Tensor): Done flags from the environment. Shape: (num_envs,) + extras (dict): Extra information from the environment. + + Observations: + + The observations TensorDict usually contains multiple observation groups. The `obs_groups` + dictionary of the runner configuration specifies which observation groups are used for which + purpose, i.e., it maps the available observation groups to observation sets. The observation sets + (keys of the `obs_groups` dictionary) currently used by rsl_rl are: + + - "policy": Specified observation groups are used as input to the actor/student network. + - "critic": Specified observation groups are used as input to the critic network. + - "teacher": Specified observation groups are used as input to the teacher network. + - "rnd_state": Specified observation groups are used as input to the RND network. + + Incomplete or incorrect configurations are handled in the `resolve_obs_groups()` function in + `rsl_rl/utils/utils.py`. + + Extras: + + The extras dictionary includes metrics such as the episode reward, episode length, etc. The following + dictionary keys are used by rsl_rl: + + - "time_outs" (torch.Tensor): Timeouts for the environments. These correspond to terminations that + happen due to time limits and not due to the environment reaching a terminal state. This is useful + for environments that have a fixed episode length. + + - "log" (dict[str, float | torch.Tensor]): Additional information for logging and debugging purposes. + The key should be a string and start with "/" for namespacing. The value can be a scalar or a + tensor. If it is a tensor, the mean of the tensor is used for logging. + """ + raise NotImplementedError + \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/__init__.py new file mode 100644 index 0000000..9afe1a2 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Definitions for neural-network components for RL-agents.""" + +from .actor_critic import ActorCritic +from .actor_critic_recurrent import ActorCriticRecurrent +from .rnd import * +from .student_teacher import StudentTeacher +from .student_teacher_recurrent import StudentTeacherRecurrent +from .symmetry import * +from .discriminator_multi import DiscriminatorMulti +__all__ = [ + "ActorCritic", + "ActorCriticRecurrent", + "StudentTeacher", + "StudentTeacherRecurrent", + "DiscriminatorMulti", +] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py new file mode 100644 index 0000000..0efc36b --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py @@ -0,0 +1,195 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +from torch.distributions import Normal + +from rsl_rl.networks import MLP, EmpiricalNormalization + + +class ActorCritic(nn.Module): + is_recurrent = False + + def __init__( + self, + obs, + obs_groups, + num_actions, + actor_obs_normalization=False, + critic_obs_normalization=False, + actor_hidden_dims=[256, 256, 256], + critic_hidden_dims=[256, 256, 256], + activation="elu", + init_noise_std=1.0, + noise_std_type: str = "scalar", + state_dependent_std=False, + **kwargs, + ): + if kwargs: + print( + "ActorCritic.__init__ got unexpected arguments, which will be ignored: " + + str([key for key in kwargs.keys()]) + ) + super().__init__() + + # get the observation dimensions + self.obs_groups = obs_groups + num_actor_obs = 0 + for obs_group in obs_groups["policy"]: + assert len(obs[obs_group].shape) == 2, "The ActorCritic module only supports 1D observations." + num_actor_obs += obs[obs_group].shape[-1] + num_critic_obs = 0 + for obs_group in obs_groups["critic"]: + assert len(obs[obs_group].shape) == 2, "The ActorCritic module only supports 1D observations." + num_critic_obs += obs[obs_group].shape[-1] + + self.state_dependent_std = state_dependent_std + # actor + if self.state_dependent_std: + self.actor = MLP(num_actor_obs, [2, num_actions], actor_hidden_dims, activation) + else: + self.actor = MLP(num_actor_obs, num_actions, actor_hidden_dims, activation) + # actor observation normalization + self.actor_obs_normalization = actor_obs_normalization + if actor_obs_normalization: + self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs) + else: + self.actor_obs_normalizer = torch.nn.Identity() + print(f"Actor MLP: {self.actor}") + + # critic + self.critic = MLP(num_critic_obs, 1, critic_hidden_dims, activation) + # critic observation normalization + self.critic_obs_normalization = critic_obs_normalization + if critic_obs_normalization: + self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs) + else: + self.critic_obs_normalizer = torch.nn.Identity() + print(f"Critic MLP: {self.critic}") + + # Action noise + self.noise_std_type = noise_std_type + if self.state_dependent_std: + torch.nn.init.zeros_(self.actor[-2].weight[num_actions:]) + if self.noise_std_type == "scalar": + torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std) + elif self.noise_std_type == "log": + torch.nn.init.constant_( + self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7)) + ) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + else: + if self.noise_std_type == "scalar": + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + elif self.noise_std_type == "log": + self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + + # Action distribution (populated in update_distribution) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args(False) + + def reset(self, dones=None): + pass + + def forward(self): + raise NotImplementedError + + @property + def action_mean(self): + return self.distribution.mean + + @property + def action_std(self): + return self.distribution.stddev + + @property + def entropy(self): + return self.distribution.entropy().sum(dim=-1) + + def update_distribution(self, obs): + if self.state_dependent_std: + # compute mean and standard deviation + mean_and_std = self.actor(obs) + if self.noise_std_type == "scalar": + mean, std = torch.unbind(mean_and_std, dim=-2) + elif self.noise_std_type == "log": + mean, log_std = torch.unbind(mean_and_std, dim=-2) + std = torch.exp(log_std) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + else: + # compute mean + mean = self.actor(obs) + # compute standard deviation + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + # create distribution + self.distribution = Normal(mean, std) + + def act(self, obs, **kwargs): + obs = self.get_actor_obs(obs) + obs = self.actor_obs_normalizer(obs) + self.update_distribution(obs) + return self.distribution.sample() + + def act_inference(self, obs): + obs = self.get_actor_obs(obs) + obs = self.actor_obs_normalizer(obs) + return self.actor(obs) + + def evaluate(self, obs, **kwargs): + obs = self.get_critic_obs(obs) + obs = self.critic_obs_normalizer(obs) + return self.critic(obs) + + def get_actor_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["policy"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_critic_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["critic"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_actions_log_prob(self, actions): + return self.distribution.log_prob(actions).sum(dim=-1) + + def update_normalization(self, obs): + if self.actor_obs_normalization: + actor_obs = self.get_actor_obs(obs) + self.actor_obs_normalizer.update(actor_obs) + if self.critic_obs_normalization: + critic_obs = self.get_critic_obs(obs) + self.critic_obs_normalizer.update(critic_obs) + + def load_state_dict(self, state_dict, strict=True): + """Load the parameters of the actor-critic model. + + Args: + state_dict (dict): State dictionary of the model. + strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this + module's state_dict() function. + + Returns: + bool: Whether this training resumes a previous training. This flag is used by the `load()` function of + `OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation). + """ + + super().load_state_dict(state_dict, strict=strict) + return True # training resumes diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py new file mode 100644 index 0000000..bba46ca --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py @@ -0,0 +1,218 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +import warnings +from torch.distributions import Normal + +from rsl_rl.networks import MLP, EmpiricalNormalization, Memory + + +class ActorCriticRecurrent(nn.Module): + is_recurrent = True + + def __init__( + self, + obs, + obs_groups, + num_actions, + actor_obs_normalization=False, + critic_obs_normalization=False, + actor_hidden_dims=[256, 256, 256], + critic_hidden_dims=[256, 256, 256], + activation="elu", + init_noise_std=1.0, + noise_std_type: str = "scalar", + state_dependent_std=False, + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=1, + **kwargs, + ): + if "rnn_hidden_size" in kwargs: + warnings.warn( + "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. " + "Please use `rnn_hidden_dim` instead.", + DeprecationWarning, + ) + if rnn_hidden_dim == 256: # Only override if the new argument is at its default + rnn_hidden_dim = kwargs.pop("rnn_hidden_size") + if kwargs: + print( + "ActorCriticRecurrent.__init__ got unexpected arguments, which will be ignored: " + str(kwargs.keys()), + ) + super().__init__() + + # get the observation dimensions + self.obs_groups = obs_groups + num_actor_obs = 0 + for obs_group in obs_groups["policy"]: + assert len(obs[obs_group].shape) == 2, "The ActorCriticRecurrent module only supports 1D observations." + num_actor_obs += obs[obs_group].shape[-1] + num_critic_obs = 0 + for obs_group in obs_groups["critic"]: + assert len(obs[obs_group].shape) == 2, "The ActorCriticRecurrent module only supports 1D observations." + num_critic_obs += obs[obs_group].shape[-1] + + self.state_dependent_std = state_dependent_std + # actor + self.memory_a = Memory(num_actor_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) + if self.state_dependent_std: + self.actor = MLP(rnn_hidden_dim, [2, num_actions], actor_hidden_dims, activation) + else: + self.actor = MLP(rnn_hidden_dim, num_actions, actor_hidden_dims, activation) + + # actor observation normalization + self.actor_obs_normalization = actor_obs_normalization + if actor_obs_normalization: + self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs) + else: + self.actor_obs_normalizer = torch.nn.Identity() + print(f"Actor RNN: {self.memory_a}") + print(f"Actor MLP: {self.actor}") + + # critic + self.memory_c = Memory(num_critic_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) + self.critic = MLP(rnn_hidden_dim, 1, critic_hidden_dims, activation) + # critic observation normalization + self.critic_obs_normalization = critic_obs_normalization + if critic_obs_normalization: + self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs) + else: + self.critic_obs_normalizer = torch.nn.Identity() + print(f"Critic RNN: {self.memory_c}") + print(f"Critic MLP: {self.critic}") + + # Action noise + self.noise_std_type = noise_std_type + if self.state_dependent_std: + torch.nn.init.zeros_(self.actor[-2].weight[num_actions:]) + if self.noise_std_type == "scalar": + torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std) + elif self.noise_std_type == "log": + torch.nn.init.constant_( + self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7)) + ) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + else: + if self.noise_std_type == "scalar": + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + elif self.noise_std_type == "log": + self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + + # Action distribution (populated in update_distribution) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args(False) + + @property + def action_mean(self): + return self.distribution.mean + + @property + def action_std(self): + return self.distribution.stddev + + @property + def entropy(self): + return self.distribution.entropy().sum(dim=-1) + + def reset(self, dones=None): + self.memory_a.reset(dones) + self.memory_c.reset(dones) + + def forward(self): + raise NotImplementedError + + def update_distribution(self, obs): + if self.state_dependent_std: + # compute mean and standard deviation + mean_and_std = self.actor(obs) + if self.noise_std_type == "scalar": + mean, std = torch.unbind(mean_and_std, dim=-2) + elif self.noise_std_type == "log": + mean, log_std = torch.unbind(mean_and_std, dim=-2) + std = torch.exp(log_std) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + else: + # compute mean + mean = self.actor(obs) + # compute standard deviation + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + # create distribution + self.distribution = Normal(mean, std) + + def act(self, obs, masks=None, hidden_states=None): + obs = self.get_actor_obs(obs) + obs = self.actor_obs_normalizer(obs) + out_mem = self.memory_a(obs, masks, hidden_states).squeeze(0) + self.update_distribution(out_mem) + return self.distribution.sample() + + def act_inference(self, obs): + obs = self.get_actor_obs(obs) + obs = self.actor_obs_normalizer(obs) + out_mem = self.memory_a(obs).squeeze(0) + return self.actor(out_mem) + + def evaluate(self, obs, masks=None, hidden_states=None): + obs = self.get_critic_obs(obs) + obs = self.critic_obs_normalizer(obs) + out_mem = self.memory_c(obs, masks, hidden_states).squeeze(0) + return self.critic(out_mem) + + def get_actor_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["policy"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_critic_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["critic"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_actions_log_prob(self, actions): + return self.distribution.log_prob(actions).sum(dim=-1) + + def get_hidden_states(self): + return self.memory_a.hidden_states, self.memory_c.hidden_states + + def update_normalization(self, obs): + if self.actor_obs_normalization: + actor_obs = self.get_actor_obs(obs) + self.actor_obs_normalizer.update(actor_obs) + if self.critic_obs_normalization: + critic_obs = self.get_critic_obs(obs) + self.critic_obs_normalizer.update(critic_obs) + + def load_state_dict(self, state_dict, strict=True): + """Load the parameters of the actor-critic model. + + Args: + state_dict (dict): State dictionary of the model. + strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this + module's state_dict() function. + + Returns: + bool: Whether this training resumes a previous training. This flag is used by the `load()` function of + `OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation). + """ + + super().load_state_dict(state_dict, strict=strict) + return True diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py new file mode 100644 index 0000000..7899d23 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py @@ -0,0 +1,102 @@ +import torch +import torch.nn as nn +from torch import autograd +import torch.nn.utils.spectral_norm as spectral_norm + + +class DiscriminatorMulti(nn.Module): + def __init__( + self, state_dim, amp_reward_coef, hidden_layer_sizes, device, + num_frames=2, task_reward_lerp=0.0, use_lerp=True): + super(DiscriminatorMulti, self).__init__() + + self.device = device + self.state_dim = state_dim + self.use_lerp = use_lerp + self.num_frames = num_frames # 存储帧数参数 + + self.amp_reward_coef = amp_reward_coef + amp_layers = [] + + curr_in_dim = state_dim * num_frames + for hidden_dim in hidden_layer_sizes: + amp_layers.append(spectral_norm(nn.Linear(curr_in_dim, hidden_dim))) + amp_layers.append(nn.ReLU()) + curr_in_dim = hidden_dim + self.trunk = nn.Sequential(*amp_layers).to(device) + self.amp_linear = spectral_norm(nn.Linear(hidden_layer_sizes[-1], 1)).to(device) + + self.trunk.train() + self.amp_linear.train() + + self.task_reward_lerp = task_reward_lerp + + def forward(self, x): + h = self.trunk(x) + d = self.amp_linear(h) + return d + + def compute_grad_pen(self, + expert_states, # 改为接收多帧状态列表 + lambda_=10): + # 将多帧状态沿最后一个维度拼接 + expert_data = expert_states.flatten(1) + expert_data.requires_grad = True + + disc = self.amp_linear(self.trunk(expert_data)) + ones = torch.ones(disc.size(), device=disc.device) + grad = autograd.grad( + outputs=disc, inputs=expert_data, + grad_outputs=ones, create_graph=True, + retain_graph=True, only_inputs=True)[0] + + # Enforce that the grad norm approaches 0. + grad_pen = lambda_ * (grad.norm(2, dim=1) - 0).pow(2).mean() + return grad_pen + + + def get_disc_weights(self): + weights = [] + for m in self.trunk.modules(): + if isinstance(m, nn.Linear): + weights.append(torch.flatten(m.weight)) + + weights.append(torch.flatten(self.amp_linear.weight)) + return weights + + def get_disc_logit_weights(self): + return torch.flatten(self.amp_linear.weight) + + def predict_amp_reward( + self, states, # 改为接收多帧状态列表 + task_reward, normalizer=None): + """ + states: torch.Tensor, shape=(num_envs, num_frames, state_dim) + task_reward: torch.Tensor, shape=(num_envs, 1) + """ + # import ipdb; ipdb.set_trace() + with torch.no_grad(): + self.eval() + if normalizer is not None: + # 对每一帧状态进行归一化 + states = normalizer.normalize_torch(states, self.device) + + # 拼接多帧状态 + state_cat = states.flatten(1) + d = self.amp_linear(self.trunk(state_cat)) + disc_reward = self.amp_reward_coef * torch.clamp(1 - (1/4) * torch.square(d - 1), min=0) + + if self.use_lerp: + if self.task_reward_lerp > 0: + reward = self._lerp_reward(disc_reward, task_reward.unsqueeze(-1)) + self.train() + return reward.squeeze(), d, disc_reward.squeeze() * (1.0 - self.task_reward_lerp) + else: + disc_reward *= 0.02 + reward = task_reward.unsqueeze(-1) + disc_reward + self.train() + return reward.squeeze(), d, disc_reward.squeeze() + + def _lerp_reward(self, disc_r, task_r): + r = (1.0 - self.task_reward_lerp) * disc_r + self.task_reward_lerp * task_r + return r \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/rnd.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/rnd.py new file mode 100644 index 0000000..8e65c43 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/rnd.py @@ -0,0 +1,209 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn + +from rsl_rl.networks import MLP, EmpiricalDiscountedVariationNormalization, EmpiricalNormalization + + +class RandomNetworkDistillation(nn.Module): + """Implementation of Random Network Distillation (RND) [1] + + References: + .. [1] Burda, Yuri, et al. "Exploration by random network distillation." arXiv preprint arXiv:1810.12894 (2018). + """ + + def __init__( + self, + num_states: int, + obs_groups: dict, + num_outputs: int, + predictor_hidden_dims: list[int], + target_hidden_dims: list[int], + activation: str = "elu", + weight: float = 0.0, + state_normalization: bool = False, + reward_normalization: bool = False, + device: str = "cpu", + weight_schedule: dict | None = None, + ): + """Initialize the RND module. + + - If :attr:`state_normalization` is True, then the input state is normalized using an Empirical Normalization layer. + - If :attr:`reward_normalization` is True, then the intrinsic reward is normalized using an Empirical Discounted + Variation Normalization layer. + + .. note:: + If the hidden dimensions are -1 in the predictor and target networks configuration, then the number of states + is used as the hidden dimension. + + Args: + num_states: Number of states/inputs to the predictor and target networks. + num_outputs: Number of outputs (embedding size) of the predictor and target networks. + predictor_hidden_dims: List of hidden dimensions of the predictor network. + target_hidden_dims: List of hidden dimensions of the target network. + activation: Activation function. Defaults to "elu". + weight: Scaling factor of the intrinsic reward. Defaults to 0.0. + state_normalization: Whether to normalize the input state. Defaults to False. + reward_normalization: Whether to normalize the intrinsic reward. Defaults to False. + device: Device to use. Defaults to "cpu". + weight_schedule: The type of schedule to use for the RND weight parameter. + Defaults to None, in which case the weight parameter is constant. + It is a dictionary with the following keys: + + - "mode": The type of schedule to use for the RND weight parameter. + - "constant": Constant weight schedule. + - "step": Step weight schedule. + - "linear": Linear weight schedule. + + For the "step" weight schedule, the following parameters are required: + + - "final_step": The step at which the weight parameter is set to the final value. + - "final_value": The final value of the weight parameter. + + For the "linear" weight schedule, the following parameters are required: + - "initial_step": The step at which the weight parameter is set to the initial value. + - "final_step": The step at which the weight parameter is set to the final value. + - "final_value": The final value of the weight parameter. + """ + # initialize parent class + super().__init__() + + # Store parameters + self.num_states = num_states + self.obs_groups = obs_groups + self.num_outputs = num_outputs + self.initial_weight = weight + self.device = device + self.state_normalization = state_normalization + self.reward_normalization = reward_normalization + + # Normalization of input gates + if state_normalization: + self.state_normalizer = EmpiricalNormalization(shape=[self.num_states], until=1.0e8).to(self.device) + else: + self.state_normalizer = torch.nn.Identity() + # Normalization of intrinsic reward + if reward_normalization: + self.reward_normalizer = EmpiricalDiscountedVariationNormalization(shape=[], until=1.0e8).to(self.device) + else: + self.reward_normalizer = torch.nn.Identity() + + # counter for the number of updates + self.update_counter = 0 + + # resolve weight schedule + if weight_schedule is not None: + self.weight_scheduler_params = weight_schedule + self.weight_scheduler = getattr(self, f"_{weight_schedule['mode']}_weight_schedule") + else: + self.weight_scheduler = None + # Create network architecture + self.predictor = MLP(num_states, num_outputs, predictor_hidden_dims, activation).to(self.device) + self.target = MLP(num_states, num_outputs, target_hidden_dims, activation).to(self.device) + + # make target network not trainable + self.target.eval() + + def get_intrinsic_reward(self, obs) -> torch.Tensor: + # Note: the counter is updated number of env steps per learning iteration + self.update_counter += 1 + # Extract the rnd state from the observation + rnd_state = self.get_rnd_state(obs) + rnd_state = self.state_normalizer(rnd_state) + # Obtain the embedding of the rnd state from the target and predictor networks + target_embedding = self.target(rnd_state).detach() + predictor_embedding = self.predictor(rnd_state).detach() + # Compute the intrinsic reward as the distance between the embeddings + intrinsic_reward = torch.linalg.norm(target_embedding - predictor_embedding, dim=1) + # Normalize intrinsic reward + intrinsic_reward = self.reward_normalizer(intrinsic_reward) + + # Check the weight schedule + if self.weight_scheduler is not None: + self.weight = self.weight_scheduler(step=self.update_counter, **self.weight_scheduler_params) + else: + self.weight = self.initial_weight + # Scale intrinsic reward + intrinsic_reward *= self.weight + + return intrinsic_reward + + def forward(self, *args, **kwargs): + raise RuntimeError("Forward method is not implemented. Use get_intrinsic_reward instead.") + + def train(self, mode: bool = True): + # sets module into training mode + self.predictor.train(mode) + if self.state_normalization: + self.state_normalizer.train(mode) + if self.reward_normalization: + self.reward_normalizer.train(mode) + return self + + def eval(self): + return self.train(False) + + def get_rnd_state(self, obs): + obs_list = [] + for obs_group in self.obs_groups["rnd_state"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def update_normalization(self, obs): + # Normalize the state + if self.state_normalization: + rnd_state = self.get_rnd_state(obs) + self.state_normalizer.update(rnd_state) + + """ + Different weight schedules. + """ + + def _constant_weight_schedule(self, step: int, **kwargs): + return self.initial_weight + + def _step_weight_schedule(self, step: int, final_step: int, final_value: float, **kwargs): + return self.initial_weight if step < final_step else final_value + + def _linear_weight_schedule(self, step: int, initial_step: int, final_step: int, final_value: float, **kwargs): + if step < initial_step: + return self.initial_weight + elif step > final_step: + return final_value + else: + return self.initial_weight + (final_value - self.initial_weight) * (step - initial_step) / ( + final_step - initial_step + ) + + +def resolve_rnd_config(alg_cfg, obs, obs_groups, env): + """Resolve the RND configuration. + + Args: + alg_cfg: The algorithm configuration dictionary. + obs: The observation dictionary. + obs_groups: The observation groups dictionary. + env: The environment. + + Returns: + The resolved algorithm configuration dictionary. + """ + # resolve dimension of rnd gated state + if "rnd_cfg" in alg_cfg and alg_cfg["rnd_cfg"] is not None: + # get dimension of rnd gated state + num_rnd_state = 0 + for obs_group in obs_groups["rnd_state"]: + assert len(obs[obs_group].shape) == 2, "The RND module only supports 1D observations." + num_rnd_state += obs[obs_group].shape[-1] + # add rnd gated state to config + alg_cfg["rnd_cfg"]["num_states"] = num_rnd_state + alg_cfg["rnd_cfg"]["obs_groups"] = obs_groups + # scale down the rnd weight with timestep + alg_cfg["rnd_cfg"]["weight"] *= env.unwrapped.step_dt + return alg_cfg diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py new file mode 100644 index 0000000..6bf1380 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py @@ -0,0 +1,206 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +from torch.distributions import Normal + +from rsl_rl.networks import MLP, EmpiricalNormalization + + +class StudentTeacher(nn.Module): + is_recurrent = False + + def __init__( + self, + obs, + obs_groups, + num_actions, + student_obs_normalization=False, + teacher_obs_normalization=False, + student_hidden_dims=[256, 256, 256], + teacher_hidden_dims=[256, 256, 256], + activation="elu", + init_noise_std=0.1, + noise_std_type: str = "scalar", + **kwargs, + ): + if kwargs: + print( + "StudentTeacher.__init__ got unexpected arguments, which will be ignored: " + + str([key for key in kwargs.keys()]) + ) + super().__init__() + + self.loaded_teacher = False # indicates if teacher has been loaded + + # get the observation dimensions + self.obs_groups = obs_groups + num_student_obs = 0 + for obs_group in obs_groups["policy"]: + assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." + num_student_obs += obs[obs_group].shape[-1] + num_teacher_obs = 0 + for obs_group in obs_groups["teacher"]: + assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." + num_teacher_obs += obs[obs_group].shape[-1] + + # student + self.student = MLP(num_student_obs, num_actions, student_hidden_dims, activation) + + # student observation normalization + self.student_obs_normalization = student_obs_normalization + if student_obs_normalization: + self.student_obs_normalizer = EmpiricalNormalization(num_student_obs) + else: + self.student_obs_normalizer = torch.nn.Identity() + + print(f"Student MLP: {self.student}") + + # teacher + self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation) + self.teacher.eval() + + # teacher observation normalization + self.teacher_obs_normalization = teacher_obs_normalization + if teacher_obs_normalization: + self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs) + else: + self.teacher_obs_normalizer = torch.nn.Identity() + + print(f"Teacher MLP: {self.teacher}") + + # action noise + self.noise_std_type = noise_std_type + if self.noise_std_type == "scalar": + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + elif self.noise_std_type == "log": + self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + + # action distribution (populated in update_distribution) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args(False) + + def reset(self, dones=None, hidden_states=None): + pass + + def forward(self): + raise NotImplementedError + + @property + def action_mean(self): + return self.distribution.mean + + @property + def action_std(self): + return self.distribution.stddev + + @property + def entropy(self): + return self.distribution.entropy().sum(dim=-1) + + def update_distribution(self, obs): + # compute mean + mean = self.student(obs) + # compute standard deviation + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + # create distribution + self.distribution = Normal(mean, std) + + def act(self, obs): + obs = self.get_student_obs(obs) + obs = self.student_obs_normalizer(obs) + self.update_distribution(obs) + return self.distribution.sample() + + def act_inference(self, obs): + obs = self.get_student_obs(obs) + obs = self.student_obs_normalizer(obs) + return self.student(obs) + + def evaluate(self, obs): + obs = self.get_teacher_obs(obs) + obs = self.teacher_obs_normalizer(obs) + with torch.no_grad(): + return self.teacher(obs) + + def get_student_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["policy"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_teacher_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["teacher"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_hidden_states(self): + return None + + def detach_hidden_states(self, dones=None): + pass + + def train(self, mode=True): + super().train(mode) + # make sure teacher is in eval mode + self.teacher.eval() + self.teacher_obs_normalizer.eval() + + def update_normalization(self, obs): + if self.student_obs_normalization: + student_obs = self.get_student_obs(obs) + self.student_obs_normalizer.update(student_obs) + + def load_state_dict(self, state_dict, strict=True): + """Load the parameters of the student and teacher networks. + + Args: + state_dict (dict): State dictionary of the model. + strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this + module's state_dict() function. + + Returns: + bool: Whether this training resumes a previous training. This flag is used by the `load()` function of + `OnPolicyRunner` to determine how to load further parameters. + """ + + # check if state_dict contains teacher and student or just teacher parameters + if any("actor" in key for key in state_dict.keys()): # loading parameters from rl training + # rename keys to match teacher and remove critic parameters + teacher_state_dict = {} + teacher_obs_normalizer_state_dict = {} + for key, value in state_dict.items(): + if "actor." in key: + teacher_state_dict[key.replace("actor.", "")] = value + if "actor_obs_normalizer." in key: + teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value + self.teacher.load_state_dict(teacher_state_dict, strict=strict) + self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict) + # set flag for successfully loading the parameters + self.loaded_teacher = True + self.teacher.eval() + self.teacher_obs_normalizer.eval() + return False # training does not resume + elif any("student" in key for key in state_dict.keys()): # loading parameters from distillation training + super().load_state_dict(state_dict, strict=strict) + # set flag for successfully loading the parameters + self.loaded_teacher = True + self.teacher.eval() + self.teacher_obs_normalizer.eval() + return True # training resumes + else: + raise ValueError("state_dict does not contain student or teacher parameters") diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py new file mode 100644 index 0000000..964a2dc --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py @@ -0,0 +1,249 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +import warnings +from torch.distributions import Normal + +from rsl_rl.networks import MLP, EmpiricalNormalization, Memory + + +class StudentTeacherRecurrent(nn.Module): + is_recurrent = True + + def __init__( + self, + obs, + obs_groups, + num_actions, + student_obs_normalization=False, + teacher_obs_normalization=False, + student_hidden_dims=[256, 256, 256], + teacher_hidden_dims=[256, 256, 256], + activation="elu", + init_noise_std=0.1, + noise_std_type: str = "scalar", + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=1, + teacher_recurrent=False, + **kwargs, + ): + if "rnn_hidden_size" in kwargs: + warnings.warn( + "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. " + "Please use `rnn_hidden_dim` instead.", + DeprecationWarning, + ) + if rnn_hidden_dim == 256: # Only override if the new argument is at its default + rnn_hidden_dim = kwargs.pop("rnn_hidden_size") + if kwargs: + print( + "StudentTeacherRecurrent.__init__ got unexpected arguments, which will be ignored: " + + str(kwargs.keys()), + ) + super().__init__() + + self.loaded_teacher = False # indicates if teacher has been loaded + self.teacher_recurrent = teacher_recurrent # indicates if teacher is recurrent too + + # get the observation dimensions + self.obs_groups = obs_groups + num_student_obs = 0 + for obs_group in obs_groups["policy"]: + assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." + num_student_obs += obs[obs_group].shape[-1] + num_teacher_obs = 0 + for obs_group in obs_groups["teacher"]: + assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." + num_teacher_obs += obs[obs_group].shape[-1] + + # student + self.memory_s = Memory(num_student_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) + self.student = MLP(rnn_hidden_dim, num_actions, student_hidden_dims, activation) + + # student observation normalization + self.student_obs_normalization = student_obs_normalization + if student_obs_normalization: + self.student_obs_normalizer = EmpiricalNormalization(num_student_obs) + else: + self.student_obs_normalizer = torch.nn.Identity() + + print(f"Student RNN: {self.memory_s}") + print(f"Student MLP: {self.student}") + + # teacher + if self.teacher_recurrent: + self.memory_t = Memory( + num_teacher_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim + ) + num_teacher_obs = rnn_hidden_dim + self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation) + + # teacher observation normalization + self.teacher_obs_normalization = teacher_obs_normalization + if teacher_obs_normalization: + self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs) + else: + self.teacher_obs_normalizer = torch.nn.Identity() + + if self.teacher_recurrent: + print(f"Teacher RNN: {self.memory_t}") + print(f"Teacher MLP: {self.teacher}") + + # action noise + self.noise_std_type = noise_std_type + if self.noise_std_type == "scalar": + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + elif self.noise_std_type == "log": + self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + + # action distribution (populated in update_distribution) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args(False) + + def reset(self, dones=None, hidden_states=None): + if hidden_states is None: + hidden_states = (None, None) + self.memory_s.reset(dones, hidden_states[0]) + if self.teacher_recurrent: + self.memory_t.reset(dones, hidden_states[1]) + + def forward(self): + raise NotImplementedError + + @property + def action_mean(self): + return self.distribution.mean + + @property + def action_std(self): + return self.distribution.stddev + + @property + def entropy(self): + return self.distribution.entropy().sum(dim=-1) + + def update_distribution(self, obs): + # compute mean + mean = self.student(obs) + # compute standard deviation + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + # create distribution + self.distribution = Normal(mean, std) + + def act(self, obs): + obs = self.get_student_obs(obs) + obs = self.student_obs_normalizer(obs) + out_mem = self.memory_s(obs).squeeze(0) + self.update_distribution(out_mem) + return self.distribution.sample() + + def act_inference(self, obs): + obs = self.get_student_obs(obs) + obs = self.student_obs_normalizer(obs) + out_mem = self.memory_s(obs).squeeze(0) + return self.student(out_mem) + + def evaluate(self, obs): + obs = self.get_teacher_obs(obs) + obs = self.teacher_obs_normalizer(obs) + with torch.no_grad(): + if self.teacher_recurrent: + self.memory_t.eval() + obs = self.memory_t(obs).squeeze(0) + return self.teacher(obs) + + def get_student_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["policy"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_teacher_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["teacher"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_hidden_states(self): + if self.teacher_recurrent: + return self.memory_s.hidden_states, self.memory_t.hidden_states + else: + return self.memory_s.hidden_states, None + + def detach_hidden_states(self, dones=None): + self.memory_s.detach_hidden_states(dones) + if self.teacher_recurrent: + self.memory_t.detach_hidden_states(dones) + + def train(self, mode=True): + super().train(mode) + # make sure teacher is in eval mode + self.teacher.eval() + self.teacher_obs_normalizer.eval() + + def update_normalization(self, obs): + if self.student_obs_normalization: + student_obs = self.get_student_obs(obs) + self.student_obs_normalizer.update(student_obs) + + def load_state_dict(self, state_dict, strict=True): + """Load the parameters of the student and teacher networks. + + Args: + state_dict (dict): State dictionary of the model. + strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this + module's state_dict() function. + + Returns: + bool: Whether this training resumes a previous training. This flag is used by the `load()` function of + `OnPolicyRunner` to determine how to load further parameters. + """ + + # check if state_dict contains teacher and student or just teacher parameters + if any("actor" in key for key in state_dict.keys()): # loading parameters from rl training + # rename keys to match teacher and remove critic parameters + teacher_state_dict = {} + teacher_obs_normalizer_state_dict = {} + for key, value in state_dict.items(): + if "actor." in key: + teacher_state_dict[key.replace("actor.", "")] = value + if "actor_obs_normalizer." in key: + teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value + self.teacher.load_state_dict(teacher_state_dict, strict=strict) + self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict) + # also load recurrent memory if teacher is recurrent + if self.teacher_recurrent: + memory_t_state_dict = {} + for key, value in state_dict.items(): + if "memory_a." in key: + memory_t_state_dict[key.replace("memory_a.", "")] = value + self.memory_t.load_state_dict(memory_t_state_dict, strict=strict) + # set flag for successfully loading the parameters + self.loaded_teacher = True + self.teacher.eval() + self.teacher_obs_normalizer.eval() + return False # training does not resume + elif any("student" in key for key in state_dict.keys()): # loading parameters from distillation training + super().load_state_dict(state_dict, strict=strict) + # set flag for successfully loading the parameters + self.loaded_teacher = True + self.teacher.eval() + self.teacher_obs_normalizer.eval() + return True # training resumes + else: + raise ValueError("state_dict does not contain student or teacher parameters") diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/symmetry.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/symmetry.py new file mode 100644 index 0000000..b017515 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/symmetry.py @@ -0,0 +1,24 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + + +def resolve_symmetry_config(alg_cfg, env): + """Resolve the symmetry configuration. + + Args: + alg_cfg: The algorithm configuration dictionary. + env: The environment. + + Returns: + The resolved algorithm configuration dictionary. + """ + + # if using symmetry then pass the environment config object + if "symmetry_cfg" in alg_cfg and alg_cfg["symmetry_cfg"] is not None: + # this is used by the symmetry function for handling different observation terms + alg_cfg["symmetry_cfg"]["_env"] = env + return alg_cfg diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/__init__.py new file mode 100644 index 0000000..c18f487 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Definitions for components of modules.""" + +from .memory import Memory +from .mlp import MLP +from .normalization import EmpiricalDiscountedVariationNormalization, EmpiricalNormalization diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/memory.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/memory.py new file mode 100644 index 0000000..7577357 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/memory.py @@ -0,0 +1,70 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch.nn as nn + +from rsl_rl.utils import unpad_trajectories + + +class Memory(nn.Module): + """Memory module for recurrent networks. + + This module is used to store the hidden states of the policy. + Currently only supports GRU and LSTM. + """ + + def __init__(self, input_size, type="lstm", num_layers=1, hidden_size=256): + super().__init__() + # RNN + rnn_cls = nn.GRU if type.lower() == "gru" else nn.LSTM + self.rnn = rnn_cls(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers) + self.hidden_states = None + + def forward(self, input, masks=None, hidden_states=None): + batch_mode = masks is not None + if batch_mode: + # batch mode: needs saved hidden states + if hidden_states is None: + raise ValueError("Hidden states not passed to memory module during policy update") + out, _ = self.rnn(input, hidden_states) + out = unpad_trajectories(out, masks) + else: + # inference/distillation mode: uses hidden states of last step + out, self.hidden_states = self.rnn(input.unsqueeze(0), self.hidden_states) + return out + + def reset(self, dones=None, hidden_states=None): + if dones is None: # reset all hidden states + if hidden_states is None: + self.hidden_states = None + else: + self.hidden_states = hidden_states + elif self.hidden_states is not None: # reset hidden states of done environments + if hidden_states is None: + if isinstance(self.hidden_states, tuple): # tuple in case of LSTM + for hidden_state in self.hidden_states: + hidden_state[..., dones == 1, :] = 0.0 + else: + self.hidden_states[..., dones == 1, :] = 0.0 + else: + NotImplementedError( + "Resetting hidden states of done environments with custom hidden states is not implemented" + ) + + def detach_hidden_states(self, dones=None): + if self.hidden_states is not None: + if dones is None: # detach all hidden states + if isinstance(self.hidden_states, tuple): # tuple in case of LSTM + self.hidden_states = tuple(hidden_state.detach() for hidden_state in self.hidden_states) + else: + self.hidden_states = self.hidden_states.detach() + else: # detach hidden states of done environments + if isinstance(self.hidden_states, tuple): # tuple in case of LSTM + for hidden_state in self.hidden_states: + hidden_state[..., dones == 1, :] = hidden_state[..., dones == 1, :].detach() + else: + self.hidden_states[..., dones == 1, :] = self.hidden_states[..., dones == 1, :].detach() diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/mlp.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/mlp.py new file mode 100644 index 0000000..e91574e --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/mlp.py @@ -0,0 +1,120 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +from functools import reduce + +from rsl_rl.utils import resolve_nn_activation + + +class MLP(nn.Sequential): + """Multi-layer perceptron. + + The MLP network is a sequence of linear layers and activation functions. The + last layer is a linear layer that outputs the desired dimension unless the + last activation function is specified. + + It provides additional conveniences: + + - If the hidden dimensions have a value of ``-1``, the dimension is inferred + from the input dimension. + - If the output dimension is a tuple, the output is reshaped to the desired + shape. + + """ + + def __init__( + self, + input_dim: int, + output_dim: int | tuple[int] | list[int], + hidden_dims: tuple[int] | list[int], + activation: str = "elu", + last_activation: str | None = None, + ): + """Initialize the MLP. + + Args: + input_dim: Dimension of the input. + output_dim: Dimension of the output. + hidden_dims: Dimensions of the hidden layers. A value of ``-1`` indicates + that the dimension should be inferred from the input dimension. + activation: Activation function. Defaults to "elu". + last_activation: Activation function of the last layer. Defaults to None, + in which case the last layer is linear. + """ + super().__init__() + + # resolve activation functions + activation_mod = resolve_nn_activation(activation) + last_activation_mod = resolve_nn_activation(last_activation) if last_activation is not None else None + # resolve number of hidden dims if they are -1 + hidden_dims_processed = [input_dim if dim == -1 else dim for dim in hidden_dims] + + # create layers sequentially + layers = [] + layers.append(nn.Linear(input_dim, hidden_dims_processed[0])) + layers.append(activation_mod) + + for layer_index in range(len(hidden_dims_processed) - 1): + layers.append(nn.Linear(hidden_dims_processed[layer_index], hidden_dims_processed[layer_index + 1])) + layers.append(activation_mod) + + # add last layer + if isinstance(output_dim, int): + layers.append(nn.Linear(hidden_dims_processed[-1], output_dim)) + else: + # compute the total output dimension + total_out_dim = reduce(lambda x, y: x * y, output_dim) + # add a layer to reshape the output to the desired shape + layers.append(nn.Linear(hidden_dims_processed[-1], total_out_dim)) + layers.append(nn.Unflatten(dim=-1, unflattened_size=output_dim)) + + # add last activation function if specified + if last_activation_mod is not None: + layers.append(last_activation_mod) + + # register the layers + for idx, layer in enumerate(layers): + self.add_module(f"{idx}", layer) + + def init_weights(self, scales: float | tuple[float]): + """Initialize the weights of the MLP. + + Args: + scales: Scale factor for the weights. + """ + + def get_scale(idx) -> float: + """Get the scale factor for the weights of the MLP. + + Args: + idx: Index of the layer. + """ + return scales[idx] if isinstance(scales, (list, tuple)) else scales + + # initialize the weights + for idx, module in enumerate(self): + if isinstance(module, nn.Linear): + nn.init.orthogonal_(module.weight, gain=get_scale(idx)) + nn.init.zeros_(module.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass of the MLP. + + Args: + x: Input tensor. + """ + for layer in self: + x = layer(x) + return x + + def reset(self, dones=None, hidden_states=None): + pass + + def detach_hidden_states(self, dones=None): + pass diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/normalization.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/normalization.py new file mode 100644 index 0000000..5fd9692 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/normalization.py @@ -0,0 +1,130 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Copyright (c) 2020 Preferred Networks, Inc. + +from __future__ import annotations + +import torch +from torch import nn + + +class EmpiricalNormalization(nn.Module): + """Normalize mean and variance of values based on empirical values.""" + + def __init__(self, shape, eps=1e-2, until=None): + """Initialize EmpiricalNormalization module. + + Args: + shape (int or tuple of int): Shape of input values except batch axis. + eps (float): Small value for stability. + until (int or None): If this arg is specified, the module learns input values until the sum of batch sizes + exceeds it. + + Note: The normalization parameters are computed over the whole batch, not for each environment separately. + """ + super().__init__() + self.eps = eps + self.until = until + self.register_buffer("_mean", torch.zeros(shape).unsqueeze(0)) + self.register_buffer("_var", torch.ones(shape).unsqueeze(0)) + self.register_buffer("_std", torch.ones(shape).unsqueeze(0)) + self.register_buffer("count", torch.tensor(0, dtype=torch.long)) + + @property + def mean(self): + return self._mean.squeeze(0).clone() + + @property + def std(self): + return self._std.squeeze(0).clone() + + def forward(self, x): + """Normalize mean and variance of values based on empirical values.""" + + return (x - self._mean) / (self._std + self.eps) + + @torch.jit.unused + def update(self, x): + """Learn input values without computing the output values of them""" + + if not self.training: + return + if self.until is not None and self.count >= self.until: + return + + count_x = x.shape[0] + self.count += count_x + rate = count_x / self.count + var_x = torch.var(x, dim=0, unbiased=False, keepdim=True) + mean_x = torch.mean(x, dim=0, keepdim=True) + delta_mean = mean_x - self._mean + self._mean += rate * delta_mean + self._var += rate * (var_x - self._var + delta_mean * (mean_x - self._mean)) + self._std = torch.sqrt(self._var) + + @torch.jit.unused + def inverse(self, y): + """De-normalize values based on empirical values.""" + + return y * (self._std + self.eps) + self._mean + + +class EmpiricalDiscountedVariationNormalization(nn.Module): + """Reward normalization from Pathak's large scale study on PPO. + + Reward normalization. Since the reward function is non-stationary, it is useful to normalize + the scale of the rewards so that the value function can learn quickly. We did this by dividing + the rewards by a running estimate of the standard deviation of the sum of discounted rewards. + """ + + def __init__(self, shape, eps=1e-2, gamma=0.99, until=None): + super().__init__() + + self.emp_norm = EmpiricalNormalization(shape, eps, until) + self.disc_avg = _DiscountedAverage(gamma) + + def forward(self, rew): + if self.training: + # update discounted rewards + avg = self.disc_avg.update(rew) + # update moments from discounted rewards + self.emp_norm.update(avg) + + # normalize rewards with the empirical std + if self.emp_norm._std > 0: + return rew / self.emp_norm._std + else: + return rew + + +""" +Helper class. +""" + + +class _DiscountedAverage: + r"""Discounted average of rewards. + + The discounted average is defined as: + + .. math:: + + \bar{R}_t = \gamma \bar{R}_{t-1} + r_t + + Args: + gamma (float): Discount factor. + """ + + def __init__(self, gamma): + self.avg = None + self.gamma = gamma + + def update(self, rew: torch.Tensor) -> torch.Tensor: + if self.avg is None: + self.avg = rew + else: + self.avg = self.avg * self.gamma + rew + return self.avg diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/__init__.py new file mode 100644 index 0000000..61f1682 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Implementation of runners for environment-agent interaction.""" + +from .on_policy_runner import OnPolicyRunner # isort:skip +from .distillation_runner import DistillationRunner +from .amp_on_policy_runner import AMPOnPolicyRunner + +__all__ = ["OnPolicyRunner", "DistillationRunner", "AMPOnPolicyRunner"] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py new file mode 100644 index 0000000..c0b9b9e --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py @@ -0,0 +1,521 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +import statistics +import time +import torch +import warnings +from collections import deque + +import rsl_rl +from rsl_rl.algorithms import AMP_PPO +from rsl_rl.env import VecEnv +from rsl_rl.modules import ActorCritic, ActorCriticRecurrent,DiscriminatorMulti, resolve_rnd_config, resolve_symmetry_config +from rsl_rl.utils import resolve_obs_groups, store_code_state, Normalizer, G1_AMPLoader + + +class AMPOnPolicyRunner: + """On-policy runner for training and evaluation of actor-critic methods.""" + + def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): + self.cfg = train_cfg + self.alg_cfg = train_cfg["algorithm"] + self.policy_cfg = train_cfg["policy"] + self.device = device + self.env = env + + # check if multi-gpu is enabled + self._configure_multi_gpu() + + # store training configuration + self.num_steps_per_env = self.cfg["num_steps_per_env"] + self.save_interval = self.cfg["save_interval"] + + # query observations from environment for algorithm construction + obs = self.env.get_observations() + default_sets = ["critic"] + if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None: + default_sets.append("rnd_state") + self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets) + + self.amp_data = G1_AMPLoader( + device, + time_between_frames=1/50.0, + preload_transitions=True, + num_preload_transitions=train_cfg["amp_num_preload_transitions"], + motion_files=train_cfg["amp_motion_files"], + num_frames=train_cfg['amp_num_frames'] + ) + + self.amp_observation_dim = self.amp_data.observation_dim if self.cfg["amp_num_obs"] == 0 else self.cfg["amp_num_obs"] + self.amp_num_frames = 0 if self.cfg["amp_num_frames"] == 0 else self.cfg["amp_num_frames"] + self.amp_normalizer = Normalizer(self.amp_observation_dim) + self.discriminator = DiscriminatorMulti( + self.amp_observation_dim, + train_cfg["amp_reward_coef"], + train_cfg["amp_discr_hidden_dims"], + device, + train_cfg["amp_num_frames"], + train_cfg["amp_task_reward_lerp"], + train_cfg['use_lerp'], + ).to(self.device) + + + # create the algorithm + self.alg = self._construct_algorithm(obs) + + # Decide whether to disable logging + # We only log from the process with rank 0 (main process) + self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 + + # Logging + self.log_dir = log_dir + self.writer = None + self.tot_timesteps = 0 + self.tot_time = 0 + self.current_learning_iteration = 0 + self.git_status_repos = [rsl_rl.__file__] + + def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 + # initialize writer + self._prepare_logging_writer() + + # randomize initial episode lengths (for exploration) + if init_at_random_ep_len: + self.env.episode_length_buf = torch.randint_like( + self.env.episode_length_buf, high=int(self.env.max_episode_length) + ) + + # start learning + obs = self.env.get_observations().to(self.device) + self.train_mode() # switch to train mode (for dropout for example) + amp_obs = self.env.get_amp_observations() + amp_obs = amp_obs.to(self.device) + if self.amp_num_frames != 0: + self.amp_obs_frames = torch.zeros(size=(self.env.num_envs, self.amp_num_frames, self.amp_observation_dim), device=self.device) + self.amp_obs_frames = torch.concat((self.amp_obs_frames[:, 1:], amp_obs.unsqueeze(1)), dim=1) + + # Book keeping + ep_infos = [] + rewbuffer = deque(maxlen=100) + lenbuffer = deque(maxlen=100) + step_discrewbuffer = deque(maxlen=100) + + cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_single_step_disc_rew = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + # create buffers for logging extrinsic and intrinsic rewards + if self.alg.rnd: + erewbuffer = deque(maxlen=100) + irewbuffer = deque(maxlen=100) + cur_ereward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_ireward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + # Ensure all parameters are in-synced + if self.is_distributed: + print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") + self.alg.broadcast_parameters() + + # Start training + start_iter = self.current_learning_iteration + tot_iter = start_iter + num_learning_iterations + for it in range(start_iter, tot_iter): + start = time.time() + # Rollout + with torch.inference_mode(): + for _ in range(self.num_steps_per_env): + # Sample actions + actions = self.alg.act(obs,amp_obs) + # Step the environment + obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) + # Move to device + obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) + + next_amp_obs = self.env.get_amp_observations() + next_amp_obs = next_amp_obs.to(self.device) + next_amp_obs_with_term = torch.clone(next_amp_obs) + + reset_env_ids = self.env.reset_env_ids + terminal_amp_states = self.env.get_amp_observations()[reset_env_ids] + next_amp_obs_with_term[reset_env_ids] = terminal_amp_states + self.amp_obs_frames = torch.concat((self.amp_obs_frames[:, 1:], next_amp_obs_with_term.unsqueeze(1)), dim=1) + + amp_reward = torch.zeros(self.env.num_envs, device=obs.device) + + mask = self.env.contact_phase[:, 0] == 1.0 + if mask.any(): + rewards[mask], logit, disc_reward = self.alg.discriminator.predict_amp_reward( + self.amp_obs_frames[mask], rewards[mask], normalizer=self.alg.amp_normalizer + ) + amp_reward[mask] += disc_reward + + # process the step + self.alg.process_env_step(obs, rewards, dones, extras, next_amp_obs_with_term, self.amp_obs_frames) + self.amp_obs_frames[reset_env_ids] = 0 + + amp_obs = torch.clone(next_amp_obs) + # Extract intrinsic rewards (only for logging) + intrinsic_rewards = self.alg.intrinsic_rewards if self.alg.rnd else None + # book keeping + if self.log_dir is not None: + if "episode" in extras: + ep_infos.append(extras["episode"]) + elif "log" in extras: + ep_infos.append(extras["log"]) + # Update rewards + if self.alg.rnd: + cur_ereward_sum += rewards + cur_ireward_sum += intrinsic_rewards # type: ignore + cur_reward_sum += rewards + intrinsic_rewards + else: + cur_reward_sum += rewards + # Update episode length + cur_episode_length += 1 + cur_single_step_disc_rew += amp_reward + # Clear data for completed episodes + # -- common + new_ids = (dones > 0).nonzero(as_tuple=False) + rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) + lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) + cur_reward_sum[new_ids] = 0 + cur_episode_length[new_ids] = 0 + to_extend_disc = (cur_single_step_disc_rew[new_ids] / self.env.max_episode_length_s)[:, 0].cpu().numpy() + step_discrewbuffer.extend(to_extend_disc.tolist()) + cur_single_step_disc_rew[new_ids] = 0 + # -- intrinsic and extrinsic rewards + if self.alg.rnd: + erewbuffer.extend(cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist()) + irewbuffer.extend(cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist()) + cur_ereward_sum[new_ids] = 0 + cur_ireward_sum[new_ids] = 0 + + stop = time.time() + collection_time = stop - start + start = stop + + # compute returns + self.alg.compute_returns(obs) + + # update policy + loss_dict = self.alg.update() + + stop = time.time() + learn_time = stop - start + self.current_learning_iteration = it + # log info + if self.log_dir is not None and not self.disable_logs: + # Log information + self.log(locals()) + # Save model + if it % self.save_interval == 0: + self.save(os.path.join(self.log_dir, f"model_{it}.pt")) + + # Clear episode infos + ep_infos.clear() + # Save code state + if it == start_iter and not self.disable_logs: + # obtain all the diff files + git_file_paths = store_code_state(self.log_dir, self.git_status_repos) + # if possible store them to wandb + if self.logger_type in ["wandb", "neptune"] and git_file_paths: + for path in git_file_paths: + self.writer.save_file(path) + + # Save the final model after training + if self.log_dir is not None and not self.disable_logs: + self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) + + def log(self, locs: dict, width: int = 80, pad: int = 35): + # Compute the collection size + collection_size = self.num_steps_per_env * self.env.num_envs * self.gpu_world_size + # Update total time-steps and time + self.tot_timesteps += collection_size + self.tot_time += locs["collection_time"] + locs["learn_time"] + iteration_time = locs["collection_time"] + locs["learn_time"] + + # -- Episode info + ep_string = "" + if locs["ep_infos"]: + for key in locs["ep_infos"][0]: + infotensor = torch.tensor([], device=self.device) + for ep_info in locs["ep_infos"]: + # handle scalar and zero dimensional tensor infos + if key not in ep_info: + continue + if not isinstance(ep_info[key], torch.Tensor): + ep_info[key] = torch.Tensor([ep_info[key]]) + if len(ep_info[key].shape) == 0: + ep_info[key] = ep_info[key].unsqueeze(0) + infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) + value = torch.mean(infotensor) + # log to logger and terminal + if "/" in key: + self.writer.add_scalar(key, value, locs["it"]) + ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + else: + self.writer.add_scalar("Episode/" + key, value, locs["it"]) + ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" + + mean_std = self.alg.policy.action_std.mean() + fps = int(collection_size / (locs["collection_time"] + locs["learn_time"])) + + # -- Losses + for key, value in locs["loss_dict"].items(): + self.writer.add_scalar(f"Loss/{key}", value, locs["it"]) + self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) + + # -- Policy + self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) + + # -- Performance + self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) + self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) + self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) + + # -- Training + if len(locs["rewbuffer"]) > 0: + # separate logging for intrinsic and extrinsic rewards + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(locs["erewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(locs["irewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/weight", self.alg.rnd.weight, locs["it"]) + # everything else + self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) + self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) + self.writer.add_scalar('Train/mean_step_disc_reward', statistics.mean(locs['step_discrewbuffer']), locs['it']) + if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging + self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) + self.writer.add_scalar( + "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time + ) + + str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " + + if len(locs["rewbuffer"]) > 0: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + f"""{'Step disc reward:':>{pad}} {statistics.mean(locs['step_discrewbuffer']):.2f}\n""" + ) + # -- Losses + for key, value in locs["loss_dict"].items(): + log_string += f"""{f'Mean {key} loss:':>{pad}} {value:.4f}\n""" + # -- Rewards + if hasattr(self.alg, "rnd") and self.alg.rnd: + log_string += ( + f"""{'Mean extrinsic reward:':>{pad}} {statistics.mean(locs['erewbuffer']):.2f}\n""" + f"""{'Mean intrinsic reward:':>{pad}} {statistics.mean(locs['irewbuffer']):.2f}\n""" + ) + log_string += f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" + # -- episode info + log_string += f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" + else: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + ) + for key, value in locs["loss_dict"].items(): + log_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + + log_string += ep_string + log_string += ( + f"""{'-' * width}\n""" + f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" + f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" + f"""{'Time elapsed:':>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n""" + f"""{'ETA:':>{pad}} {time.strftime( + "%H:%M:%S", + time.gmtime( + self.tot_time / (locs['it'] - locs['start_iter'] + 1) + * (locs['start_iter'] + locs['num_learning_iterations'] - locs['it']) + ) + )}\n""" + ) + print(log_string) + + def save(self, path: str, infos=None): + # -- Save model + saved_dict = { + "model_state_dict": self.alg.policy.state_dict(), + "optimizer_state_dict": self.alg.optimizer.state_dict(), + "iter": self.current_learning_iteration, + "infos": infos, + } + # -- Save RND model if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict() + saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict() + torch.save(saved_dict, path) + + # upload model to external logging service + if self.logger_type in ["neptune", "wandb"] and not self.disable_logs: + self.writer.save_model(path, self.current_learning_iteration) + + def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None): + loaded_dict = torch.load(path, weights_only=False, map_location=map_location) + # -- Load model + resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"]) + # -- Load RND model if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"]) + # -- load optimizer if used + if load_optimizer and resumed_training: + # -- algorithm optimizer + self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) + # -- RND optimizer if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"]) + # -- load current learning iteration + if resumed_training: + self.current_learning_iteration = loaded_dict["iter"] + return loaded_dict["infos"] + + def get_inference_policy(self, device=None): + self.eval_mode() # switch to evaluation mode (dropout for example) + if device is not None: + self.alg.policy.to(device) + return self.alg.policy.act_inference + + def train_mode(self): + # -- PPO + self.alg.policy.train() + self.alg.discriminator.train() + # -- RND + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.train() + + def eval_mode(self): + # -- PPO + self.alg.policy.eval() + self.alg.discriminator.eval() + # -- RND + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.eval() + + def add_git_repo_to_log(self, repo_file_path): + self.git_status_repos.append(repo_file_path) + + """ + Helper functions. + """ + + def _configure_multi_gpu(self): + """Configure multi-gpu training.""" + # check if distributed training is enabled + self.gpu_world_size = int(os.getenv("WORLD_SIZE", "1")) + self.is_distributed = self.gpu_world_size > 1 + + # if not distributed training, set local and global rank to 0 and return + if not self.is_distributed: + self.gpu_local_rank = 0 + self.gpu_global_rank = 0 + self.multi_gpu_cfg = None + return + + # get rank and world size + self.gpu_local_rank = int(os.getenv("LOCAL_RANK", "0")) + self.gpu_global_rank = int(os.getenv("RANK", "0")) + + # make a configuration dictionary + self.multi_gpu_cfg = { + "global_rank": self.gpu_global_rank, # rank of the main process + "local_rank": self.gpu_local_rank, # rank of the current process + "world_size": self.gpu_world_size, # total number of processes + } + + # check if user has device specified for local rank + if self.device != f"cuda:{self.gpu_local_rank}": + raise ValueError( + f"Device '{self.device}' does not match expected device for local rank '{self.gpu_local_rank}'." + ) + # validate multi-gpu configuration + if self.gpu_local_rank >= self.gpu_world_size: + raise ValueError( + f"Local rank '{self.gpu_local_rank}' is greater than or equal to world size '{self.gpu_world_size}'." + ) + if self.gpu_global_rank >= self.gpu_world_size: + raise ValueError( + f"Global rank '{self.gpu_global_rank}' is greater than or equal to world size '{self.gpu_world_size}'." + ) + + # initialize torch distributed + torch.distributed.init_process_group(backend="nccl", rank=self.gpu_global_rank, world_size=self.gpu_world_size) + # set device to the local rank + torch.cuda.set_device(self.gpu_local_rank) + + def _construct_algorithm(self, obs) -> AMP_PPO: + """Construct the actor-critic algorithm.""" + # resolve RND config + self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env) + + # resolve symmetry config + self.alg_cfg = resolve_symmetry_config(self.alg_cfg, self.env) + + # resolve deprecated normalization config + if self.cfg.get("empirical_normalization") is not None: + warnings.warn( + "The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and " + "`critic_obs_normalization` as part of the `policy` configuration instead.", + DeprecationWarning, + ) + if self.policy_cfg.get("actor_obs_normalization") is None: + self.policy_cfg["actor_obs_normalization"] = self.cfg["empirical_normalization"] + if self.policy_cfg.get("critic_obs_normalization") is None: + self.policy_cfg["critic_obs_normalization"] = self.cfg["empirical_normalization"] + + # initialize the actor-critic + actor_critic_class = eval(self.policy_cfg.pop("class_name")) + actor_critic: ActorCritic | ActorCriticRecurrent = actor_critic_class( + obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg + ).to(self.device) + + # initialize the algorithm + alg_class = eval(self.alg_cfg.pop("class_name")) + + alg: AMP_PPO = alg_class(actor_critic, self.discriminator, self.amp_data, self.amp_normalizer, self.amp_num_frames, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg) + + # initialize the storage + alg.init_storage( + "rl", + self.env.num_envs, + self.num_steps_per_env, + obs, + [self.env.num_actions], + ) + + return alg + + def _prepare_logging_writer(self): + """Prepares the logging writers.""" + if self.log_dir is not None and self.writer is None and not self.disable_logs: + # Launch either Tensorboard or Neptune & Tensorboard summary writer(s), default: Tensorboard. + self.logger_type = self.cfg.get("logger", "tensorboard") + self.logger_type = self.logger_type.lower() + + if self.logger_type == "neptune": + from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter + + self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "wandb": + from rsl_rl.utils.wandb_utils import WandbSummaryWriter + + self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "tensorboard": + from torch.utils.tensorboard import SummaryWriter + + self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) + else: + raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.") diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py new file mode 100644 index 0000000..9cc6a8b --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py @@ -0,0 +1,179 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +import time +import torch +from collections import deque + +import rsl_rl +from rsl_rl.algorithms import Distillation +from rsl_rl.env import VecEnv +from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent +from rsl_rl.runners import OnPolicyRunner +from rsl_rl.utils import resolve_obs_groups, store_code_state + + +class DistillationRunner(OnPolicyRunner): + """On-policy runner for training and evaluation of teacher-student training.""" + + def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): + self.cfg = train_cfg + self.alg_cfg = train_cfg["algorithm"] + self.policy_cfg = train_cfg["policy"] + self.device = device + self.env = env + + # check if multi-gpu is enabled + self._configure_multi_gpu() + + # store training configuration + self.num_steps_per_env = self.cfg["num_steps_per_env"] + self.save_interval = self.cfg["save_interval"] + + # query observations from environment for algorithm construction + obs = self.env.get_observations() + self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets=["teacher"]) + + # create the algorithm + self.alg = self._construct_algorithm(obs) + + # Decide whether to disable logging + # We only log from the process with rank 0 (main process) + self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 + + # Logging + self.log_dir = log_dir + self.writer = None + self.tot_timesteps = 0 + self.tot_time = 0 + self.current_learning_iteration = 0 + self.git_status_repos = [rsl_rl.__file__] + + def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 + # initialize writer + self._prepare_logging_writer() + # check if teacher is loaded + if not self.alg.policy.loaded_teacher: + raise ValueError("Teacher model parameters not loaded. Please load a teacher model to distill.") + + # randomize initial episode lengths (for exploration) + if init_at_random_ep_len: + self.env.episode_length_buf = torch.randint_like( + self.env.episode_length_buf, high=int(self.env.max_episode_length) + ) + + # start learning + obs = self.env.get_observations().to(self.device) + self.train_mode() # switch to train mode (for dropout for example) + + # Book keeping + ep_infos = [] + rewbuffer = deque(maxlen=100) + lenbuffer = deque(maxlen=100) + cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + # Ensure all parameters are in-synced + if self.is_distributed: + print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") + self.alg.broadcast_parameters() + + # Start training + start_iter = self.current_learning_iteration + tot_iter = start_iter + num_learning_iterations + for it in range(start_iter, tot_iter): + start = time.time() + # Rollout + with torch.inference_mode(): + for _ in range(self.num_steps_per_env): + # Sample actions + actions = self.alg.act(obs) + # Step the environment + obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) + # Move to device + obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) + # process the step + self.alg.process_env_step(obs, rewards, dones, extras) + # book keeping + if self.log_dir is not None: + if "episode" in extras: + ep_infos.append(extras["episode"]) + elif "log" in extras: + ep_infos.append(extras["log"]) + # Update rewards + cur_reward_sum += rewards + # Update episode length + cur_episode_length += 1 + # Clear data for completed episodes + new_ids = (dones > 0).nonzero(as_tuple=False) + rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) + lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) + cur_reward_sum[new_ids] = 0 + cur_episode_length[new_ids] = 0 + + stop = time.time() + collection_time = stop - start + start = stop + + # update policy + loss_dict = self.alg.update() + + stop = time.time() + learn_time = stop - start + self.current_learning_iteration = it + # log info + if self.log_dir is not None and not self.disable_logs: + # Log information + self.log(locals()) + # Save model + if it % self.save_interval == 0: + self.save(os.path.join(self.log_dir, f"model_{it}.pt")) + + # Clear episode infos + ep_infos.clear() + # Save code state + if it == start_iter and not self.disable_logs: + # obtain all the diff files + git_file_paths = store_code_state(self.log_dir, self.git_status_repos) + # if possible store them to wandb + if self.logger_type in ["wandb", "neptune"] and git_file_paths: + for path in git_file_paths: + self.writer.save_file(path) + + # Save the final model after training + if self.log_dir is not None and not self.disable_logs: + self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) + + """ + Helper methods. + """ + + def _construct_algorithm(self, obs) -> Distillation: + """Construct the distillation algorithm.""" + # initialize the actor-critic + student_teacher_class = eval(self.policy_cfg.pop("class_name")) + student_teacher: StudentTeacher | StudentTeacherRecurrent = student_teacher_class( + obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg + ).to(self.device) + + # initialize the algorithm + alg_class = eval(self.alg_cfg.pop("class_name")) + alg: Distillation = alg_class( + student_teacher, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg + ) + + # initialize the storage + alg.init_storage( + "distillation", + self.env.num_envs, + self.num_steps_per_env, + obs, + [self.env.num_actions], + ) + + return alg diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py new file mode 100644 index 0000000..36f11f3 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py @@ -0,0 +1,460 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +import statistics +import time +import torch +import warnings +from collections import deque + +import rsl_rl +from rsl_rl.algorithms import PPO +from rsl_rl.env import VecEnv +from rsl_rl.modules import ActorCritic, ActorCriticRecurrent, resolve_rnd_config, resolve_symmetry_config +from rsl_rl.utils import resolve_obs_groups, store_code_state + + +class OnPolicyRunner: + """On-policy runner for training and evaluation of actor-critic methods.""" + + def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): + self.cfg = train_cfg + self.alg_cfg = train_cfg["algorithm"] + self.policy_cfg = train_cfg["policy"] + self.device = device + self.env = env + + # check if multi-gpu is enabled + self._configure_multi_gpu() + + # store training configuration + self.num_steps_per_env = self.cfg["num_steps_per_env"] + self.save_interval = self.cfg["save_interval"] + + # query observations from environment for algorithm construction + obs = self.env.get_observations() + default_sets = ["critic"] + if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None: + default_sets.append("rnd_state") + self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets) + + # create the algorithm + self.alg = self._construct_algorithm(obs) + + # Decide whether to disable logging + # We only log from the process with rank 0 (main process) + self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 + + # Logging + self.log_dir = log_dir + self.writer = None + self.tot_timesteps = 0 + self.tot_time = 0 + self.current_learning_iteration = 0 + self.git_status_repos = [rsl_rl.__file__] + + def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 + # initialize writer + self._prepare_logging_writer() + + # randomize initial episode lengths (for exploration) + if init_at_random_ep_len: + self.env.episode_length_buf = torch.randint_like( + self.env.episode_length_buf, high=int(self.env.max_episode_length) + ) + + # start learning + obs = self.env.get_observations().to(self.device) + self.train_mode() # switch to train mode (for dropout for example) + + # Book keeping + ep_infos = [] + rewbuffer = deque(maxlen=100) + lenbuffer = deque(maxlen=100) + cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + # create buffers for logging extrinsic and intrinsic rewards + if self.alg.rnd: + erewbuffer = deque(maxlen=100) + irewbuffer = deque(maxlen=100) + cur_ereward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_ireward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + # Ensure all parameters are in-synced + if self.is_distributed: + print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") + self.alg.broadcast_parameters() + + # Start training + start_iter = self.current_learning_iteration + tot_iter = start_iter + num_learning_iterations + for it in range(start_iter, tot_iter): + start = time.time() + # Rollout + with torch.inference_mode(): + for _ in range(self.num_steps_per_env): + # Sample actions + actions = self.alg.act(obs) + # Step the environment + obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) + # Move to device + obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) + # process the step + self.alg.process_env_step(obs, rewards, dones, extras) + # Extract intrinsic rewards (only for logging) + intrinsic_rewards = self.alg.intrinsic_rewards if self.alg.rnd else None + # book keeping + if self.log_dir is not None: + if "episode" in extras: + ep_infos.append(extras["episode"]) + elif "log" in extras: + ep_infos.append(extras["log"]) + # Update rewards + if self.alg.rnd: + cur_ereward_sum += rewards + cur_ireward_sum += intrinsic_rewards # type: ignore + cur_reward_sum += rewards + intrinsic_rewards + else: + cur_reward_sum += rewards + # Update episode length + cur_episode_length += 1 + # Clear data for completed episodes + # -- common + new_ids = (dones > 0).nonzero(as_tuple=False) + rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) + lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) + cur_reward_sum[new_ids] = 0 + cur_episode_length[new_ids] = 0 + # -- intrinsic and extrinsic rewards + if self.alg.rnd: + erewbuffer.extend(cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist()) + irewbuffer.extend(cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist()) + cur_ereward_sum[new_ids] = 0 + cur_ireward_sum[new_ids] = 0 + + stop = time.time() + collection_time = stop - start + start = stop + + # compute returns + self.alg.compute_returns(obs) + + # update policy + loss_dict = self.alg.update() + + stop = time.time() + learn_time = stop - start + self.current_learning_iteration = it + # log info + if self.log_dir is not None and not self.disable_logs: + # Log information + self.log(locals()) + # Save model + if it % self.save_interval == 0: + self.save(os.path.join(self.log_dir, f"model_{it}.pt")) + + # Clear episode infos + ep_infos.clear() + # Save code state + if it == start_iter and not self.disable_logs: + # obtain all the diff files + git_file_paths = store_code_state(self.log_dir, self.git_status_repos) + # if possible store them to wandb + if self.logger_type in ["wandb", "neptune"] and git_file_paths: + for path in git_file_paths: + self.writer.save_file(path) + + # Save the final model after training + if self.log_dir is not None and not self.disable_logs: + self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) + + def log(self, locs: dict, width: int = 80, pad: int = 35): + # Compute the collection size + collection_size = self.num_steps_per_env * self.env.num_envs * self.gpu_world_size + # Update total time-steps and time + self.tot_timesteps += collection_size + self.tot_time += locs["collection_time"] + locs["learn_time"] + iteration_time = locs["collection_time"] + locs["learn_time"] + + # -- Episode info + ep_string = "" + if locs["ep_infos"]: + for key in locs["ep_infos"][0]: + infotensor = torch.tensor([], device=self.device) + for ep_info in locs["ep_infos"]: + # handle scalar and zero dimensional tensor infos + if key not in ep_info: + continue + if not isinstance(ep_info[key], torch.Tensor): + ep_info[key] = torch.Tensor([ep_info[key]]) + if len(ep_info[key].shape) == 0: + ep_info[key] = ep_info[key].unsqueeze(0) + infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) + value = torch.mean(infotensor) + # log to logger and terminal + if "/" in key: + self.writer.add_scalar(key, value, locs["it"]) + ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + else: + self.writer.add_scalar("Episode/" + key, value, locs["it"]) + ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" + + mean_std = self.alg.policy.action_std.mean() + fps = int(collection_size / (locs["collection_time"] + locs["learn_time"])) + + # -- Losses + for key, value in locs["loss_dict"].items(): + self.writer.add_scalar(f"Loss/{key}", value, locs["it"]) + self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) + + # -- Policy + self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) + + # -- Performance + self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) + self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) + self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) + + # -- Training + if len(locs["rewbuffer"]) > 0: + # separate logging for intrinsic and extrinsic rewards + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(locs["erewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(locs["irewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/weight", self.alg.rnd.weight, locs["it"]) + # everything else + self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) + self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) + if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging + self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) + self.writer.add_scalar( + "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time + ) + + str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " + + if len(locs["rewbuffer"]) > 0: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + ) + # -- Losses + for key, value in locs["loss_dict"].items(): + log_string += f"""{f'Mean {key} loss:':>{pad}} {value:.4f}\n""" + # -- Rewards + if hasattr(self.alg, "rnd") and self.alg.rnd: + log_string += ( + f"""{'Mean extrinsic reward:':>{pad}} {statistics.mean(locs['erewbuffer']):.2f}\n""" + f"""{'Mean intrinsic reward:':>{pad}} {statistics.mean(locs['irewbuffer']):.2f}\n""" + ) + log_string += f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" + # -- episode info + log_string += f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" + else: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + ) + for key, value in locs["loss_dict"].items(): + log_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + + log_string += ep_string + log_string += ( + f"""{'-' * width}\n""" + f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" + f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" + f"""{'Time elapsed:':>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n""" + f"""{'ETA:':>{pad}} {time.strftime( + "%H:%M:%S", + time.gmtime( + self.tot_time / (locs['it'] - locs['start_iter'] + 1) + * (locs['start_iter'] + locs['num_learning_iterations'] - locs['it']) + ) + )}\n""" + ) + print(log_string) + + def save(self, path: str, infos=None): + # -- Save model + saved_dict = { + "model_state_dict": self.alg.policy.state_dict(), + "optimizer_state_dict": self.alg.optimizer.state_dict(), + "iter": self.current_learning_iteration, + "infos": infos, + } + # -- Save RND model if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict() + saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict() + torch.save(saved_dict, path) + + # upload model to external logging service + if self.logger_type in ["neptune", "wandb"] and not self.disable_logs: + self.writer.save_model(path, self.current_learning_iteration) + + def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None): + loaded_dict = torch.load(path, weights_only=False, map_location=map_location) + # -- Load model + resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"]) + # -- Load RND model if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"]) + # -- load optimizer if used + if load_optimizer and resumed_training: + # -- algorithm optimizer + self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) + # -- RND optimizer if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"]) + # -- load current learning iteration + if resumed_training: + self.current_learning_iteration = loaded_dict["iter"] + return loaded_dict["infos"] + + def get_inference_policy(self, device=None): + self.eval_mode() # switch to evaluation mode (dropout for example) + if device is not None: + self.alg.policy.to(device) + return self.alg.policy.act_inference + + def train_mode(self): + # -- PPO + self.alg.policy.train() + # -- RND + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.train() + + def eval_mode(self): + # -- PPO + self.alg.policy.eval() + # -- RND + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.eval() + + def add_git_repo_to_log(self, repo_file_path): + self.git_status_repos.append(repo_file_path) + + """ + Helper functions. + """ + + def _configure_multi_gpu(self): + """Configure multi-gpu training.""" + # check if distributed training is enabled + self.gpu_world_size = int(os.getenv("WORLD_SIZE", "1")) + self.is_distributed = self.gpu_world_size > 1 + + # if not distributed training, set local and global rank to 0 and return + if not self.is_distributed: + self.gpu_local_rank = 0 + self.gpu_global_rank = 0 + self.multi_gpu_cfg = None + return + + # get rank and world size + self.gpu_local_rank = int(os.getenv("LOCAL_RANK", "0")) + self.gpu_global_rank = int(os.getenv("RANK", "0")) + + # make a configuration dictionary + self.multi_gpu_cfg = { + "global_rank": self.gpu_global_rank, # rank of the main process + "local_rank": self.gpu_local_rank, # rank of the current process + "world_size": self.gpu_world_size, # total number of processes + } + + # check if user has device specified for local rank + if self.device != f"cuda:{self.gpu_local_rank}": + raise ValueError( + f"Device '{self.device}' does not match expected device for local rank '{self.gpu_local_rank}'." + ) + # validate multi-gpu configuration + if self.gpu_local_rank >= self.gpu_world_size: + raise ValueError( + f"Local rank '{self.gpu_local_rank}' is greater than or equal to world size '{self.gpu_world_size}'." + ) + if self.gpu_global_rank >= self.gpu_world_size: + raise ValueError( + f"Global rank '{self.gpu_global_rank}' is greater than or equal to world size '{self.gpu_world_size}'." + ) + + # initialize torch distributed + torch.distributed.init_process_group(backend="nccl", rank=self.gpu_global_rank, world_size=self.gpu_world_size) + # set device to the local rank + torch.cuda.set_device(self.gpu_local_rank) + + def _construct_algorithm(self, obs) -> PPO: + """Construct the actor-critic algorithm.""" + # resolve RND config + self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env) + + # resolve symmetry config + self.alg_cfg = resolve_symmetry_config(self.alg_cfg, self.env) + + # resolve deprecated normalization config + if self.cfg.get("empirical_normalization") is not None: + warnings.warn( + "The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and " + "`critic_obs_normalization` as part of the `policy` configuration instead.", + DeprecationWarning, + ) + if self.policy_cfg.get("actor_obs_normalization") is None: + self.policy_cfg["actor_obs_normalization"] = self.cfg["empirical_normalization"] + if self.policy_cfg.get("critic_obs_normalization") is None: + self.policy_cfg["critic_obs_normalization"] = self.cfg["empirical_normalization"] + + # initialize the actor-critic + actor_critic_class = eval(self.policy_cfg.pop("class_name")) + actor_critic: ActorCritic | ActorCriticRecurrent = actor_critic_class( + obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg + ).to(self.device) + + # initialize the algorithm + alg_class = eval(self.alg_cfg.pop("class_name")) + alg: PPO = alg_class(actor_critic, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg) + + # initialize the storage + alg.init_storage( + "rl", + self.env.num_envs, + self.num_steps_per_env, + obs, + [self.env.num_actions], + ) + + return alg + + def _prepare_logging_writer(self): + """Prepares the logging writers.""" + if self.log_dir is not None and self.writer is None and not self.disable_logs: + # Launch either Tensorboard or Neptune & Tensorboard summary writer(s), default: Tensorboard. + self.logger_type = self.cfg.get("logger", "tensorboard") + self.logger_type = self.logger_type.lower() + + if self.logger_type == "neptune": + from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter + + self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "wandb": + from rsl_rl.utils.wandb_utils import WandbSummaryWriter + + self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "tensorboard": + from torch.utils.tensorboard import SummaryWriter + + self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) + else: + raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.") diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/__init__.py new file mode 100644 index 0000000..1624330 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Implementation of transitions storage for RL-agent.""" + +from .rollout_storage import RolloutStorage +from .replay_buffer_multi import ReplayBufferMulti +__all__ = ["RolloutStorage", "ReplayBufferMulti"] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py new file mode 100644 index 0000000..6462b8e --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py @@ -0,0 +1,38 @@ +import torch +import numpy as np + + +class ReplayBufferMulti: + """Fixed-size buffer to store experience tuples.""" + + def __init__(self, obs_dim, buffer_size, num_amp_frames, device): + """Initialize a ReplayBuffer object. + Arguments: + buffer_size (int): maximum size of buffer + """ + self.states = torch.zeros(buffer_size, num_amp_frames, obs_dim).to(device) + self.num_amp_frames = num_amp_frames + self.buffer_size = buffer_size + self.device = device + + self.step = 0 + self.num_samples = 0 + + def insert(self, states): + """Add new states to memory.""" + num_states = states.shape[0] + start_idx = self.step + end_idx = self.step + num_states + if end_idx > self.buffer_size: + self.states[self.step:self.buffer_size] = states[:self.buffer_size - self.step] + self.states[:end_idx - self.buffer_size] = states[self.buffer_size - self.step:] + else: + self.states[start_idx:end_idx] = states + + self.num_samples = min(self.buffer_size, max(end_idx, self.num_samples)) + self.step = (self.step + num_states) % self.buffer_size + + def feed_forward_generator(self, num_mini_batch, mini_batch_size): + for _ in range(num_mini_batch): + sample_idxs = np.random.choice(self.num_samples, size=mini_batch_size) + yield (self.states[sample_idxs].to(self.device)) diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py new file mode 100644 index 0000000..e9309b3 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py @@ -0,0 +1,260 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from tensordict import TensorDict + +from rsl_rl.utils import split_and_pad_trajectories + + +class RolloutStorage: + class Transition: + def __init__(self): + self.observations = None + self.actions = None + self.privileged_actions = None + self.rewards = None + self.dones = None + self.values = None + self.actions_log_prob = None + self.action_mean = None + self.action_sigma = None + self.hidden_states = None + + def clear(self): + self.__init__() + + def __init__( + self, + training_type, + num_envs, + num_transitions_per_env, + obs, + actions_shape, + device="cpu", + ): + # store inputs + self.training_type = training_type + self.device = device + self.num_transitions_per_env = num_transitions_per_env + self.num_envs = num_envs + self.actions_shape = actions_shape + + # Core + self.observations = TensorDict( + {key: torch.zeros(num_transitions_per_env, *value.shape, device=device) for key, value in obs.items()}, + batch_size=[num_transitions_per_env, num_envs], + device=self.device, + ) + self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte() + + # for distillation + if training_type == "distillation": + self.privileged_actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + + # for reinforcement learning + if training_type == "rl": + self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + + # For RNN networks + self.saved_hidden_states_a = None + self.saved_hidden_states_c = None + + # counter for the number of transitions stored + self.step = 0 + + def add_transitions(self, transition: Transition): + # check if the transition is valid + if self.step >= self.num_transitions_per_env: + raise OverflowError("Rollout buffer overflow! You should call clear() before adding new transitions.") + + # Core + self.observations[self.step].copy_(transition.observations) + self.actions[self.step].copy_(transition.actions) + self.rewards[self.step].copy_(transition.rewards.view(-1, 1)) + self.dones[self.step].copy_(transition.dones.view(-1, 1)) + + # for distillation + if self.training_type == "distillation": + self.privileged_actions[self.step].copy_(transition.privileged_actions) + + # for reinforcement learning + if self.training_type == "rl": + self.values[self.step].copy_(transition.values) + self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1)) + self.mu[self.step].copy_(transition.action_mean) + self.sigma[self.step].copy_(transition.action_sigma) + + # For RNN networks + self._save_hidden_states(transition.hidden_states) + + # increment the counter + self.step += 1 + + def _save_hidden_states(self, hidden_states): + if hidden_states is None or hidden_states == (None, None): + return + # make a tuple out of GRU hidden state sto match the LSTM format + hid_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],) + hid_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],) + # initialize if needed + if self.saved_hidden_states_a is None: + self.saved_hidden_states_a = [ + torch.zeros(self.observations.shape[0], *hid_a[i].shape, device=self.device) for i in range(len(hid_a)) + ] + self.saved_hidden_states_c = [ + torch.zeros(self.observations.shape[0], *hid_c[i].shape, device=self.device) for i in range(len(hid_c)) + ] + # copy the states + for i in range(len(hid_a)): + self.saved_hidden_states_a[i][self.step].copy_(hid_a[i]) + self.saved_hidden_states_c[i][self.step].copy_(hid_c[i]) + + def clear(self): + self.step = 0 + + def compute_returns(self, last_values, gamma, lam, normalize_advantage: bool = True): + advantage = 0 + for step in reversed(range(self.num_transitions_per_env)): + # if we are at the last step, bootstrap the return value + if step == self.num_transitions_per_env - 1: + next_values = last_values + else: + next_values = self.values[step + 1] + # 1 if we are not in a terminal state, 0 otherwise + next_is_not_terminal = 1.0 - self.dones[step].float() + # TD error: r_t + gamma * V(s_{t+1}) - V(s_t) + delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step] + # Advantage: A(s_t, a_t) = delta_t + gamma * lambda * A(s_{t+1}, a_{t+1}) + advantage = delta + next_is_not_terminal * gamma * lam * advantage + # Return: R_t = A(s_t, a_t) + V(s_t) + self.returns[step] = advantage + self.values[step] + + # Compute the advantages + self.advantages = self.returns - self.values + # Normalize the advantages if flag is set + # This is to prevent double normalization (i.e. if per minibatch normalization is used) + if normalize_advantage: + self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8) + + # for distillation + def generator(self): + if self.training_type != "distillation": + raise ValueError("This function is only available for distillation training.") + + for i in range(self.num_transitions_per_env): + yield self.observations[i], self.actions[i], self.privileged_actions[i], self.dones[i] + + # for reinforcement learning with feedforward networks + def mini_batch_generator(self, num_mini_batches, num_epochs=8): + if self.training_type != "rl": + raise ValueError("This function is only available for reinforcement learning training.") + batch_size = self.num_envs * self.num_transitions_per_env + mini_batch_size = batch_size // num_mini_batches + indices = torch.randperm(num_mini_batches * mini_batch_size, requires_grad=False, device=self.device) + + # Core + observations = self.observations.flatten(0, 1) + actions = self.actions.flatten(0, 1) + values = self.values.flatten(0, 1) + returns = self.returns.flatten(0, 1) + + # For PPO + old_actions_log_prob = self.actions_log_prob.flatten(0, 1) + advantages = self.advantages.flatten(0, 1) + old_mu = self.mu.flatten(0, 1) + old_sigma = self.sigma.flatten(0, 1) + + for epoch in range(num_epochs): + for i in range(num_mini_batches): + # Select the indices for the mini-batch + start = i * mini_batch_size + end = (i + 1) * mini_batch_size + batch_idx = indices[start:end] + + # Create the mini-batch + # -- Core + obs_batch = observations[batch_idx] + actions_batch = actions[batch_idx] + + # -- For PPO + target_values_batch = values[batch_idx] + returns_batch = returns[batch_idx] + old_actions_log_prob_batch = old_actions_log_prob[batch_idx] + advantages_batch = advantages[batch_idx] + old_mu_batch = old_mu[batch_idx] + old_sigma_batch = old_sigma[batch_idx] + + # yield the mini-batch + yield obs_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( + None, + None, + ), None + + # for reinfrocement learning with recurrent networks + def recurrent_mini_batch_generator(self, num_mini_batches, num_epochs=8): + if self.training_type != "rl": + raise ValueError("This function is only available for reinforcement learning training.") + padded_obs_trajectories, trajectory_masks = split_and_pad_trajectories(self.observations, self.dones) + + mini_batch_size = self.num_envs // num_mini_batches + for ep in range(num_epochs): + first_traj = 0 + for i in range(num_mini_batches): + start = i * mini_batch_size + stop = (i + 1) * mini_batch_size + + dones = self.dones.squeeze(-1) + last_was_done = torch.zeros_like(dones, dtype=torch.bool) + last_was_done[1:] = dones[:-1] + last_was_done[0] = True + trajectories_batch_size = torch.sum(last_was_done[:, start:stop]) + last_traj = first_traj + trajectories_batch_size + + masks_batch = trajectory_masks[:, first_traj:last_traj] + obs_batch = padded_obs_trajectories[:, first_traj:last_traj] + actions_batch = self.actions[:, start:stop] + old_mu_batch = self.mu[:, start:stop] + old_sigma_batch = self.sigma[:, start:stop] + returns_batch = self.returns[:, start:stop] + advantages_batch = self.advantages[:, start:stop] + values_batch = self.values[:, start:stop] + old_actions_log_prob_batch = self.actions_log_prob[:, start:stop] + + # reshape to [num_envs, time, num layers, hidden dim] (original shape: [time, num_layers, num_envs, hidden_dim]) + # then take only time steps after dones (flattens num envs and time dimensions), + # take a batch of trajectories and finally reshape back to [num_layers, batch, hidden_dim] + last_was_done = last_was_done.permute(1, 0) + hid_a_batch = [ + saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] + .transpose(1, 0) + .contiguous() + for saved_hidden_states in self.saved_hidden_states_a + ] + hid_c_batch = [ + saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] + .transpose(1, 0) + .contiguous() + for saved_hidden_states in self.saved_hidden_states_c + ] + # remove the tuple for GRU + hid_a_batch = hid_a_batch[0] if len(hid_a_batch) == 1 else hid_a_batch + hid_c_batch = hid_c_batch[0] if len(hid_c_batch) == 1 else hid_c_batch + + yield obs_batch, actions_batch, values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( + hid_a_batch, + hid_c_batch, + ), masks_batch + + first_traj = last_traj diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/__init__.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/__init__.py new file mode 100644 index 0000000..f5781f1 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Helper functions.""" + +from .utils import * +from .motion_loader_g1 import G1_AMPLoader + +__all__ = [ + "G1_AMPLoader", +] \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py new file mode 100644 index 0000000..677c630 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py @@ -0,0 +1,388 @@ +import os +from os.path import join as pjoin +import glob +import json +import logging + +import torch +import numpy as np +from pybullet_utils import transformations + +from rsl_rl.utils import motion_util + +_EPS = np.finfo(float).eps * 4.0 +def quaternion_slerp(q0, q1, fraction, spin=0, shortestpath=True): + """Batch quaternion spherical linear interpolation.""" + + out = torch.zeros_like(q0) + + zero_mask = torch.isclose(fraction, torch.zeros_like(fraction)).squeeze() + ones_mask = torch.isclose(fraction, torch.ones_like(fraction)).squeeze() + out[zero_mask] = q0[zero_mask] + out[ones_mask] = q1[ones_mask] + + d = torch.sum(q0 * q1, dim=-1, keepdim=True) + dist_mask = (torch.abs(torch.abs(d) - 1.0) < _EPS).squeeze() + out[dist_mask] = q0[dist_mask] + + if shortestpath: + d_old = torch.clone(d) + d = torch.where(d_old < 0, -d, d) + q1 = torch.where(d_old < 0, -q1, q1) + + angle = torch.acos(d) + spin * torch.pi + angle_mask = (torch.abs(angle) < _EPS).squeeze() + out[angle_mask] = q0[angle_mask] + + final_mask = torch.logical_or(zero_mask, ones_mask) + final_mask = torch.logical_or(final_mask, dist_mask) + final_mask = torch.logical_or(final_mask, angle_mask) + final_mask = torch.logical_not(final_mask) + + isin = 1.0 / angle + q0 *= torch.sin((1.0 - fraction) * angle) * isin + q1 *= torch.sin(fraction * angle) * isin + q0 += q1 + out[final_mask] = q0[final_mask] + return out + + +class G1_AMPLoader: + + def __init__( + self, + device, + time_between_frames, + motion_files, + preload_transitions=False, + num_preload_transitions=1000000, + num_frames=5, + ): + """Expert dataset provides AMP observations from Dog mocap dataset. + + time_between_frames: Amount of time in seconds between transition. + """ + self.device = device + self.time_between_frames = time_between_frames + self.num_frames = num_frames + + # Values to store for each trajectory. + self.trajectories = [] + self.trajectories_full = [] + self.trajectory_names = [] + self.trajectory_idxs = [] + self.trajectory_lens = [] # Traj length in seconds. + self.trajectory_weights = [] + self.trajectory_frame_durations = [] + self.trajectory_num_frames = [] + self.motion_dir = motion_files + # import ipdb; ipdb.set_trace() + for i, motion_file in enumerate(os.listdir(motion_files)): + self.trajectory_names.append(motion_file) + motion_path = pjoin(motion_files, motion_file) + motion_data = np.load(motion_path, allow_pickle=True) + motion_data_processed = np.zeros((motion_data.shape[0],36)) + + for f_i in range(motion_data.shape[0]): + motion_data_processed[f_i, :3] = motion_data[f_i, :3] # base pos + motion_data_processed[f_i, 3:7] = motion_data[f_i, 3:7] # base quat (wxyz) + motion_data_processed[f_i, 7:35] = motion_data[f_i, 7:35] # base vel + ''' + NOTE The order of motion_data_processed is + base pos 0:3, + base quat 3:7, wxyz + dof pos 7:36, (mujoco joint order) + ''' + self.trajectories.append(torch.tensor( + motion_data_processed[:, 7:], + dtype=torch.float32, + device=self.device + )) + + self.trajectories_full.append(torch.tensor( + motion_data_processed, + dtype=torch.float32, + device=self.device + )) + + self.trajectory_idxs.append(i) + self.trajectory_weights.append(1 / len(os.listdir(motion_files))) + frame_duration = 1 / 50 + + self.trajectory_frame_durations.append(frame_duration) + traj_len = (motion_data_processed.shape[0] - 1) * frame_duration # seconds + self.trajectory_lens.append(traj_len) + self.trajectory_num_frames.append(float(motion_data_processed.shape[0])) + print(f"Loaded {traj_len}s. motion from {motion_file}.") + + # Trajectory weights are used to sample some trajectories more than others. + self.trajectory_weights = np.array(self.trajectory_weights) / np.sum(self.trajectory_weights) + self.trajectory_frame_durations = np.array(self.trajectory_frame_durations) + self.trajectory_lens = np.array(self.trajectory_lens) + self.trajectory_num_frames = np.array(self.trajectory_num_frames) + + # Preload transitions. + self.preload_transitions = preload_transitions + if self.preload_transitions: + print(f'Preloading {num_preload_transitions} transitions') + + traj_idxs = self.weighted_traj_idx_sample_batch(num_preload_transitions) + times = self.traj_time_sample_batch(traj_idxs) + self.preloaded_s_prior = self.get_full_frame_at_time_batch(traj_idxs, times - self.time_between_frames) + self.preloaded_s = self.get_full_frame_at_time_batch(traj_idxs, times) + self.preloaded_s_next = self.get_full_frame_at_time_batch(traj_idxs, times + self.time_between_frames) + print(f'Finished preloading') + + # 预加载多帧数据 + self.preloaded_frames = [] + for i in range(self.num_frames): + frame_time = times + (i - (self.num_frames - 2)) * self.time_between_frames + full_frame = self.get_full_frame_at_time_batch(traj_idxs, frame_time) + # 预处理：提前提取并连接需要的列（7:26 和 29:33），避免每次生成时重复切片 + processed_frame = torch.cat([ + full_frame[:, 7:26], + full_frame[:, 29:33] + ], dim=-1) + self.preloaded_frames.append(processed_frame) + print(f'Finished preloading multiple frames') + + self.all_trajectories_full = torch.vstack(self.trajectories_full) + + def weighted_traj_idx_sample(self): + """Get traj idx via weighted sampling.""" + return np.random.choice( + self.trajectory_idxs, p=self.trajectory_weights) + + def weighted_traj_idx_sample_batch(self, size): + """Batch sample traj idxs.""" + return np.random.choice( + self.trajectory_idxs, size=size, p=self.trajectory_weights, + replace=True) + + def traj_time_sample(self, traj_idx): + """Sample random time for traj.""" + subst = self.time_between_frames + self.trajectory_frame_durations[traj_idx] + return max( + 0, (self.trajectory_lens[traj_idx] * np.random.uniform() - subst)) + + def traj_time_sample_batch(self, traj_idxs): + """Sample random time for multiple trajectories.""" + subst = self.time_between_frames + self.trajectory_frame_durations[traj_idxs] + time_samples = self.trajectory_lens[traj_idxs] * np.random.uniform(size=len(traj_idxs)) - subst + return np.maximum(np.zeros_like(time_samples), time_samples) + + def slerp(self, val0, val1, blend): + return (1.0 - blend) * val0 + blend * val1 + + def get_trajectory(self, traj_idx): + """Returns trajectory of AMP observations.""" + return self.trajectories_full[traj_idx] + + def get_frame_at_time(self, traj_idx, time): + """Returns frame for the given trajectory at the specified time.""" + p = float(time) / self.trajectory_lens[traj_idx] + n = self.trajectories[traj_idx].shape[0] + idx_low, idx_high = int(np.floor(p * n)), int(np.ceil(p * n)) + frame_start = self.trajectories[traj_idx][idx_low] + frame_end = self.trajectories[traj_idx][idx_high] + blend = p * n - idx_low + return self.slerp(frame_start, frame_end, blend) + + def get_frame_at_time_batch(self, traj_idxs, times): + """Returns frame for the given trajectory at the specified time.""" + p = times / self.trajectory_lens[traj_idxs] + n = self.trajectory_num_frames[traj_idxs] + idx_low, idx_high = np.floor(p * n).astype(np.int32), np.ceil(p * n).astype(np.int32) + all_frame_starts = torch.zeros(len(traj_idxs), self.observation_dim, device=self.device) + all_frame_ends = torch.zeros(len(traj_idxs), self.observation_dim, device=self.device) + for traj_idx in set(traj_idxs): + trajectory = self.trajectories[traj_idx] + traj_mask = traj_idxs == traj_idx + all_frame_starts[traj_mask] = trajectory[idx_low[traj_mask]] + all_frame_ends[traj_mask] = trajectory[idx_high[traj_mask]] + blend = torch.tensor(p * n - idx_low, device=self.device, dtype=torch.float32).unsqueeze(-1) + return self.slerp(all_frame_starts, all_frame_ends, blend) + + def get_full_frame_at_time(self, traj_idx, time): + """Returns full frame for the given trajectory at the specified time.""" + p = float(time) / self.trajectory_lens[traj_idx] + n = self.trajectories_full[traj_idx].shape[0] + idx_low, idx_high = int(np.floor(p * n)), int(np.ceil(p * n)) + frame_start = self.trajectories_full[traj_idx][idx_low] + frame_end = self.trajectories_full[traj_idx][idx_high] + blend = p * n - idx_low + print(idx_low, idx_high) + return self.blend_frame_pose(frame_start, frame_end, blend) + + def get_full_frame_at_time_batch(self, traj_idxs, times): + p = times / self.trajectory_lens[traj_idxs] + n = self.trajectory_num_frames[traj_idxs] + idx_low, idx_high = np.floor(p * n).astype(np.int32), np.ceil(p * n).astype(np.int32) + all_frame_pos_starts = torch.zeros(len(traj_idxs), 3, device=self.device) + all_frame_pos_ends = torch.zeros(len(traj_idxs), 3, device=self.device) + all_frame_rot_starts = torch.zeros(len(traj_idxs), 4, device=self.device) + all_frame_rot_ends = torch.zeros(len(traj_idxs), 4, device=self.device) + all_frame_amp_starts = torch.zeros(len(traj_idxs), 29, device=self.device) + all_frame_amp_ends = torch.zeros(len(traj_idxs), 29, device=self.device) + for traj_idx in set(traj_idxs): + trajectory = self.trajectories_full[traj_idx] + traj_mask = traj_idxs == traj_idx + all_frame_pos_starts[traj_mask] = G1_AMPLoader.get_root_pos_batch(trajectory[idx_low[traj_mask]]) + all_frame_pos_ends[traj_mask] = G1_AMPLoader.get_root_pos_batch(trajectory[idx_high[traj_mask]]) + all_frame_rot_starts[traj_mask] = G1_AMPLoader.get_root_rot_batch(trajectory[idx_low[traj_mask]]) + all_frame_rot_ends[traj_mask] = G1_AMPLoader.get_root_rot_batch(trajectory[idx_high[traj_mask]]) + all_frame_amp_starts[traj_mask] = trajectory[idx_low[traj_mask]][:, 7:36] # base vel3+ang3, dof vel23+ang23 + all_frame_amp_ends[traj_mask] = trajectory[idx_high[traj_mask]][:, 7:36] # base vel3+ang3, dof vel23+ang23 + blend = torch.tensor(p * n - idx_low, device=self.device, dtype=torch.float32).unsqueeze(-1) + pos_blend = self.slerp(all_frame_pos_starts, all_frame_pos_ends, blend) + rot_blend = quaternion_slerp(all_frame_rot_starts, all_frame_rot_ends, blend) + amp_blend = self.slerp(all_frame_amp_starts, all_frame_amp_ends, blend) + return torch.cat([pos_blend, rot_blend, amp_blend], dim=-1) + + def get_frame(self): + """Returns random frame.""" + traj_idx = self.weighted_traj_idx_sample() + sampled_time = self.traj_time_sample(traj_idx) + return self.get_frame_at_time(traj_idx, sampled_time) + + def get_full_frame(self): + """Returns random full frame.""" + traj_idx = self.weighted_traj_idx_sample() + sampled_time = self.traj_time_sample(traj_idx) + return self.get_full_frame_at_time(traj_idx, sampled_time) + + def get_full_frame_batch(self, num_frames): + if self.preload_transitions: + idxs = np.random.choice( + self.preloaded_s.shape[0], size=num_frames) + return self.preloaded_s[idxs] + else: + traj_idxs = self.weighted_traj_idx_sample_batch(num_frames) + times = self.traj_time_sample_batch(traj_idxs) + return self.get_full_frame_at_time_batch(traj_idxs, times) + + def blend_frame_pose(self, frame0, frame1, blend): + """Linearly interpolate between two frames, including orientation. + + Args: + frame0: First frame to be blended corresponds to (blend = 0). + frame1: Second frame to be blended corresponds to (blend = 1). + blend: Float between [0, 1], specifying the interpolation between + the two frames. + Returns: + An interpolation of the two frames. + """ + root_pos0, root_pos1 = G1_AMPLoader.get_root_pos(frame0), G1_AMPLoader.get_root_pos(frame1) + root_rot0, root_rot1 = G1_AMPLoader.get_root_rot(frame0), G1_AMPLoader.get_root_rot(frame1) + joints0, joints1 = G1_AMPLoader.get_joint_pose(frame0), G1_AMPLoader.get_joint_pose(frame1) + # tar_toe_pos_0, tar_toe_pos_1 = G1_AMPLoader.get_tar_toe_pos_local(frame0), G1_AMPLoader.get_tar_toe_pos_local(frame1) + linear_vel_0, linear_vel_1 = G1_AMPLoader.get_linear_vel(frame0), G1_AMPLoader.get_linear_vel(frame1) + angular_vel_0, angular_vel_1 = G1_AMPLoader.get_angular_vel(frame0), G1_AMPLoader.get_angular_vel(frame1) + joint_vel_0, joint_vel_1 = G1_AMPLoader.get_joint_vel(frame0), G1_AMPLoader.get_joint_vel(frame1) + + blend_root_pos = self.slerp(root_pos0, root_pos1, blend) + blend_root_rot = transformations.quaternion_slerp(root_rot0.cpu().numpy(), root_rot1.cpu().numpy(), blend) + blend_root_rot = torch.tensor(motion_util.standardize_quaternion(blend_root_rot),dtype=torch.float32, device=self.device) + blend_joints = self.slerp(joints0, joints1, blend) + # blend_tar_toe_pos = self.slerp(tar_toe_pos_0, tar_toe_pos_1, blend) + blend_linear_vel = self.slerp(linear_vel_0, linear_vel_1, blend) + blend_angular_vel = self.slerp(angular_vel_0, angular_vel_1, blend) + blend_joints_vel = self.slerp(joint_vel_0, joint_vel_1, blend) + + # return + # torch.cat([ + # blend_root_pos, blend_root_rot, blend_linear_vel, blend_angular_vel, blend_joints, blend_joints_vel]) + return torch.cat([blend_root_pos, blend_root_rot, blend_linear_vel, blend_angular_vel, blend_joints]) + + def feed_forward_generator_23dof_multi(self, num_mini_batch, mini_batch_size): + """Generates a batch of AMP transitions.""" + # import ipdb; ipdb.set_trace() + for _ in range(num_mini_batch): + if self.preload_transitions: + idxs = np.random.choice(self.preloaded_s.shape[0], size=mini_batch_size) + + frames = [] + for i in range(self.num_frames): + # 数据已在预加载时预处理，直接索引即可 + s = self.preloaded_frames[i][idxs] + frames.append(s) + else: + NotImplementedError('preload transition') + yield torch.stack(frames, dim=1) # [batch, num_frames, 16] + + + + + def quaternion_to_euler_array(self, quat): + # Ensure quaternion is in the correct format [x, y, z, w] + x, y, z, w =quat + + # Roll (x-axis rotation) + t0 = +2.0 * (w * x + y * z) + t1 = +1.0 - 2.0 * (x * x + y * y) + roll_x = np.arctan2(t0, t1) + + # Pitch (y-axis rotation) + t2 = +2.0 * (w * y - z * x) + t2 = np.clip(t2, -1.0, 1.0) + pitch_y = np.arcsin(t2) + + # Yaw (z-axis rotation) + t3 = +2.0 * (w * z + x * y) + t4 = +1.0 - 2.0 * (y * y + z * z) + yaw_z = np.arctan2(t3, t4) + + # Returns roll, pitch, yaw in a NumPy array in radians + return np.array([roll_x, pitch_y, yaw_z]) + + def euler_to_quaternion(self, root_rot): + roll, pitch, yaw = root_rot[0], root_rot[1], root_rot[2] + cy = np.cos(yaw * 0.5) + sy = np.sin(yaw * 0.5) + cp = np.cos(pitch * 0.5) + sp = np.sin(pitch * 0.5) + cr = np.cos(roll * 0.5) + sr = np.sin(roll * 0.5) + + qw = cy * cp * cr + sy * sp * sr + qx = cy * cp * sr - sy * sp * cr + qy = sy * cp * sr + cy * sp * cr + qz = sy * cp * cr - cy * sp * sr + + return np.array([qx, qy, qz, qw]) + + @property + def observation_dim(self): + """Size of AMP observations.""" + return self.trajectories[0].shape[1] + 1 + + @property + def num_motions(self): + return len(self.trajectory_names) + @staticmethod + def get_root_pos(pose): + return pose[0:3] + + @staticmethod + def get_root_pos_batch(poses): + return poses[:, 0:3] + + @staticmethod + def get_root_rot(pose): + return pose[3:7] + + @staticmethod + def get_root_rot_batch(poses): + return poses[:, 3:7] + + @staticmethod + def get_joint_pose_batch_12dof(poses): + return poses[:, 13:25] + + @staticmethod + def get_tar_toe_pos_local(pose): + return pose[G1_AMPLoader.TAR_TOE_POS_LOCAL_START_IDX:G1_AMPLoader.TAR_TOE_POS_LOCAL_END_IDX] + + @staticmethod + def get_tar_toe_pos_local_batch(poses): + return poses[:, G1_AMPLoader.TAR_TOE_POS_LOCAL_START_IDX:G1_AMPLoader.TAR_TOE_POS_LOCAL_END_IDX] diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_util.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_util.py new file mode 100644 index 0000000..3d49bfa --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_util.py @@ -0,0 +1,97 @@ +# coding=utf-8 +# Copyright 2020 The Google Research Authors. +# +# 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. + +"""Utility functions for processing motion clips.""" + +import os +import inspect +currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) +parentdir = os.path.dirname(os.path.dirname(currentdir)) +os.sys.path.insert(0, parentdir) + +import numpy as np + +from rsl_rl.utils import pose3d +# from pybullet_utils import transformations + + +def standardize_quaternion(q): + """Returns a quaternion where q.w >= 0 to remove redundancy due to q = -q. + + Args: + q: A quaternion to be standardized. + + Returns: + A quaternion with q.w >= 0. + + """ + if q[-1] < 0: + q = -q + return q + + +def normalize_rotation_angle(theta): + """Returns a rotation angle normalized between [-pi, pi]. + + Args: + theta: angle of rotation (radians). + + Returns: + An angle of rotation normalized between [-pi, pi]. + + """ + norm_theta = theta + if np.abs(norm_theta) > np.pi: + norm_theta = np.fmod(norm_theta, 2 * np.pi) + if norm_theta >= 0: + norm_theta += -2 * np.pi + else: + norm_theta += 2 * np.pi + + return norm_theta + + +def calc_heading(q): + """Returns the heading of a rotation q, specified as a quaternion. + + The heading represents the rotational component of q along the vertical + axis (z axis). + + Args: + q: A quaternion that the heading is to be computed from. + + Returns: + An angle representing the rotation about the z axis. + + """ + ref_dir = np.array([1, 0, 0]) + rot_dir = pose3d.QuaternionRotatePoint(ref_dir, q) + heading = np.arctan2(rot_dir[1], rot_dir[0]) + return heading + + +# def calc_heading_rot(q): +# """Return a quaternion representing the heading rotation of q along the vertical axis (z axis). + +# Args: +# q: A quaternion that the heading is to be computed from. + +# Returns: +# A quaternion representing the rotation about the z axis. + +# """ +# heading = calc_heading(q) +# q_heading = transformations.quaternion_about_axis(heading, [0, 0, 1]) +# return q_heading diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py new file mode 100644 index 0000000..3796ec8 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py @@ -0,0 +1,94 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +from dataclasses import asdict +from torch.utils.tensorboard import SummaryWriter + +try: + import neptune +except ModuleNotFoundError: + raise ModuleNotFoundError("neptune-client is required to log to Neptune.") + + +class NeptuneLogger: + def __init__(self, project, token): + self.run = neptune.init_run(project=project, api_token=token) + + def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): + self.run["runner_cfg"] = runner_cfg + self.run["policy_cfg"] = policy_cfg + self.run["alg_cfg"] = alg_cfg + self.run["env_cfg"] = asdict(env_cfg) + + +class NeptuneSummaryWriter(SummaryWriter): + """Summary writer for Neptune.""" + + def __init__(self, log_dir: str, flush_secs: int, cfg): + super().__init__(log_dir, flush_secs) + + try: + project = cfg["neptune_project"] + except KeyError: + raise KeyError("Please specify neptune_project in the runner config, e.g. legged_gym.") + + try: + token = os.environ["NEPTUNE_API_TOKEN"] + except KeyError: + raise KeyError( + "Neptune api token not found. Please run or add to ~/.bashrc: export NEPTUNE_API_TOKEN=YOUR_API_TOKEN" + ) + + try: + entity = os.environ["NEPTUNE_USERNAME"] + except KeyError: + raise KeyError( + "Neptune username not found. Please run or add to ~/.bashrc: export NEPTUNE_USERNAME=YOUR_USERNAME" + ) + + neptune_project = entity + "/" + project + + self.neptune_logger = NeptuneLogger(neptune_project, token) + + self.name_map = { + "Train/mean_reward/time": "Train/mean_reward_time", + "Train/mean_episode_length/time": "Train/mean_episode_length_time", + } + + run_name = os.path.split(log_dir)[-1] + + self.neptune_logger.run["log_dir"].log(run_name) + + def _map_path(self, path): + if path in self.name_map: + return self.name_map[path] + else: + return path + + def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_style=False): + super().add_scalar( + tag, + scalar_value, + global_step=global_step, + walltime=walltime, + new_style=new_style, + ) + self.neptune_logger.run[self._map_path(tag)].log(scalar_value, step=global_step) + + def stop(self): + self.neptune_logger.run.stop() + + def log_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): + self.neptune_logger.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg) + + def save_model(self, model_path, iter): + self.neptune_logger.run["model/saved_model_" + str(iter)].upload(model_path) + + def save_file(self, path, iter=None): + name = path.rsplit("/", 1)[-1].split(".")[0] + self.neptune_logger.run["git_diff/" + name].upload(path) diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/pose3d.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/pose3d.py new file mode 100644 index 0000000..a4cba1e --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/pose3d.py @@ -0,0 +1,283 @@ +# coding=utf-8 +# Copyright 2020 The Google Research Authors. +# +# 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. +"""Utilities for 3D pose conversion.""" +import math +import numpy as np + +# from pybullet_utils import transformations + +VECTOR3_0 = np.zeros(3, dtype=np.float64) +VECTOR3_1 = np.ones(3, dtype=np.float64) +VECTOR3_X = np.array([1, 0, 0], dtype=np.float64) +VECTOR3_Y = np.array([0, 1, 0], dtype=np.float64) +VECTOR3_Z = np.array([0, 0, 1], dtype=np.float64) + +# QUATERNION_IDENTITY is the multiplicative identity 1.0 + 0i + 0j + 0k. +# When interpreted as a rotation, it is the identity rotation. +QUATERNION_IDENTITY = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float64) + + +def Vector3RandomNormal(sigma, mu=VECTOR3_0): + """Returns a random 3D vector from a normal distribution. + + Each component is selected independently from a normal distribution. + + Args: + sigma: Scale (or stddev) of distribution for all variables. + mu: Mean of distribution for each variable. + + Returns: + A 3D vector in a numpy array. + """ + + random_v3 = np.random.normal(scale=sigma, size=3) + mu + return random_v3 + + +def Vector3RandomUniform(low=VECTOR3_0, high=VECTOR3_1): + """Returns a 3D vector selected uniformly from the input box. + + Args: + low: The min-value corner of the box. + high: The max-value corner of the box. + + Returns: + A 3D vector in a numpy array. + """ + + random_x = np.random.uniform(low=low[0], high=high[0]) + random_y = np.random.uniform(low=low[1], high=high[1]) + random_z = np.random.uniform(low=low[2], high=high[2]) + return np.array([random_x, random_y, random_z]) + + +def Vector3RandomUnit(): + """Returns a random 3D vector with unit length. + + Generates a 3D vector selected uniformly from the unit sphere. + + Returns: + A normalized 3D vector in a numpy array. + """ + longitude = np.random.uniform(low=-math.pi, high=math.pi) + sin_latitude = np.random.uniform(low=-1.0, high=1.0) + cos_latitude = math.sqrt(1.0 - sin_latitude * sin_latitude) + x = math.cos(longitude) * cos_latitude + y = math.sin(longitude) * cos_latitude + z = sin_latitude + return np.array([x, y, z], dtype=np.float64) + + +def QuaternionNormalize(q): + """Normalizes the quaternion to length 1. + + Divides the quaternion by its magnitude. If the magnitude is too + small, returns the quaternion identity value (1.0). + + Args: + q: A quaternion to be normalized. + + Raises: + ValueError: If input quaternion has length near zero. + + Returns: + A quaternion with magnitude 1 in a numpy array [x, y, z, w]. + + """ + q_norm = np.linalg.norm(q) + if np.isclose(q_norm, 0.0): + raise ValueError( + 'Quaternion may not be zero in QuaternionNormalize: |q| = %f, q = %s' % + (q_norm, q)) + return q / q_norm + + +def QuaternionFromAxisAngle(axis, angle): + """Returns a quaternion that generates the given axis-angle rotation. + + Returns the quaternion: sin(angle/2) * axis + cos(angle/2). + + Args: + axis: Axis of rotation, a 3D vector in a numpy array. + angle: The angle of rotation (radians). + + Raises: + ValueError: If input axis is not a normalizable 3D vector. + + Returns: + A unit quaternion in a numpy array. + + """ + if len(axis) != 3: + raise ValueError('Axis vector should have three components: %s' % axis) + axis_norm = np.linalg.norm(axis) + if np.isclose(axis_norm, 0.0): + raise ValueError('Axis vector may not have zero length: |v| = %f, v = %s' % + (axis_norm, axis)) + half_angle = angle * 0.5 + q = np.zeros(4, dtype=np.float64) + q[0:3] = axis + q[0:3] *= math.sin(half_angle) / axis_norm + q[3] = math.cos(half_angle) + return q + + +def QuaternionToAxisAngle(quat, default_axis=VECTOR3_Z, direction_axis=None): + """Calculates axis and angle of rotation performed by a quaternion. + + Calculates the axis and angle of the rotation performed by the quaternion. + The quaternion should have four values and be normalized. + + Args: + quat: Unit quaternion in a numpy array. + default_axis: 3D vector axis used if the rotation is near to zero. Without + this default, small rotations would result in an exception. It is + reasonable to use a default axis for tiny rotations, because zero angle + rotations about any axis are equivalent. + direction_axis: Used to disambiguate rotation directions. If the + direction_axis is specified, the axis of the rotation will be chosen such + that its inner product with the direction_axis is non-negative. + + Raises: + ValueError: If quat is not a normalized quaternion. + + Returns: + axis: Axis of rotation. + angle: Angle in radians. + """ + if len(quat) != 4: + raise ValueError( + 'Quaternion should have four components [x, y, z, w]: %s' % quat) + if not np.isclose(1.0, np.linalg.norm(quat)): + raise ValueError('Quaternion should have unit length: |q| = %f, q = %s' % + (np.linalg.norm(quat), quat)) + axis = quat[:3].copy() + axis_norm = np.linalg.norm(axis) + min_axis_norm = 1e-8 + if axis_norm < min_axis_norm: + axis = default_axis + if len(default_axis) != 3: + raise ValueError('Axis vector should have three components: %s' % axis) + if not np.isclose(np.linalg.norm(axis), 1.0): + raise ValueError('Axis vector should have unit length: |v| = %f, v = %s' % + (np.linalg.norm(axis), axis)) + else: + axis /= axis_norm + sin_half_angle = axis_norm + if direction_axis is not None and np.inner(axis, direction_axis) < 0: + sin_half_angle = -sin_half_angle + axis = -axis + cos_half_angle = quat[3] + half_angle = math.atan2(sin_half_angle, cos_half_angle) + angle = half_angle * 2 + return axis, angle + + +def QuaternionRandomRotation(max_angle=math.pi): + """Creates a random small rotation around a random axis. + + Generates a small rotation with the axis vector selected uniformly + from the unit sphere and an angle selected from a uniform + distribution over [0, max_angle]. + + If the max_angle is not specified, the rotation should be selected + uniformly over all possible rotation angles. + + Args: + max_angle: The maximum angle of rotation (radians). + + Returns: + A unit quaternion in a numpy array. + + """ + + angle = np.random.uniform(low=0, high=max_angle) + axis = Vector3RandomUnit() + return QuaternionFromAxisAngle(axis, angle) + + +# def QuaternionRotatePoint(point, quat): +# """Performs a rotation by quaternion. + +# Rotate the point by the quaternion using quaternion multiplication, +# (q * p * q^-1), without constructing the rotation matrix. + +# Args: +# point: The point to be rotated. +# quat: The rotation represented as a quaternion [x, y, z, w]. + +# Returns: +# A 3D vector in a numpy array. +# """ + +# q_point = np.array([point[0], point[1], point[2], 0.0]) +# quat_inverse = transformations.quaternion_inverse(quat) +# q_point_rotated = transformations.quaternion_multiply( +# transformations.quaternion_multiply(quat, q_point), quat_inverse) +# return q_point_rotated[:3] + + +def IsRotationMatrix(m): + """Returns true if the 3x3 submatrix represents a rotation. + + Args: + m: A transformation matrix. + + Raises: + ValueError: If input is not a matrix of size at least 3x3. + + Returns: + True if the 3x3 submatrix is a rotation (orthogonal). + """ + if len(m.shape) != 2 or m.shape[0] < 3 or m.shape[1] < 3: + raise ValueError('Matrix should be 3x3 or 4x4: %s\n %s' % (m.shape, m)) + rot = m[:3, :3] + eye = np.matmul(rot, np.transpose(rot)) + return np.isclose(eye, np.identity(3), atol=1e-4).all() + +# def ZAxisAlignedRobotPoseTool(robot_pose_tool): +# """Returns the current gripper pose rotated for alignment with the z-axis. + +# Args: +# robot_pose_tool: a pose3d.Pose3d() instance. + +# Returns: +# An instance of pose.Transform representing the current gripper pose +# rotated for alignment with the z-axis. +# """ +# # Align the current pose to the z-axis. +# robot_pose_tool.quaternion = transformations.quaternion_multiply( +# RotationBetween( +# robot_pose_tool.matrix4x4[0:3, 0:3].dot(np.array([0, 0, 1])), +# np.array([0.0, 0.0, -1.0])), robot_pose_tool.quaternion) +# return robot_pose_tool + +# def RotationBetween(a_translation_b, a_translation_c): +# """Computes the rotation from one vector to another. + +# The computed rotation has the property that: + +# a_translation_c = a_rotation_b_to_c * a_translation_b + +# Args: +# a_translation_b: vec3, vector to rotate from +# a_translation_c: vec3, vector to rotate to + +# Returns: +# a_rotation_b_to_c: new Orientation +# """ +# rotation = rotation3.Rotation3.rotation_between( +# a_translation_b, a_translation_c, err_msg='RotationBetween') +# return rotation.quaternion.xyzw diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/utils.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/utils.py new file mode 100644 index 0000000..3605622 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/utils.py @@ -0,0 +1,360 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import git +import importlib +import os +import pathlib +import torch +import warnings +from tensordict import TensorDict +from typing import Callable +import numpy as np +class RunningMeanStd: + def __init__(self, epsilon: float = 1e-4, shape: Tuple[int, ...] = ()): + """ + Calculates the running mean and std of a data stream + https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm + :param epsilon: helps with arithmetic issues + :param shape: the shape of the data stream's output + """ + self.mean = np.zeros(shape, np.float64) + self.var = np.ones(shape, np.float64) + self.count = epsilon + + def update(self, arr: np.ndarray) -> None: + batch_mean = np.mean(arr, axis=0) + batch_var = np.var(arr, axis=0) + batch_count = arr.shape[0] + self.update_from_moments(batch_mean, batch_var, batch_count) + + def update_from_moments(self, batch_mean: np.ndarray, batch_var: np.ndarray, batch_count: int) -> None: + delta = batch_mean - self.mean + tot_count = self.count + batch_count + + new_mean = self.mean + delta * batch_count / tot_count + m_a = self.var * self.count + m_b = batch_var * batch_count + m_2 = m_a + m_b + np.square(delta) * self.count * batch_count / (self.count + batch_count) + new_var = m_2 / (self.count + batch_count) + + new_count = batch_count + self.count + + self.mean = new_mean + self.var = new_var + self.count = new_count + + +class Normalizer(RunningMeanStd): + def __init__(self, input_dim, epsilon=1e-4, clip_obs=10.0): + super().__init__(shape=input_dim) + self.epsilon = epsilon + self.clip_obs = clip_obs + + def normalize(self, input): + return np.clip((input - self.mean) / np.sqrt(self.var + self.epsilon), -self.clip_obs, self.clip_obs) + + def normalize_torch(self, input, device): + mean_torch = torch.tensor(self.mean, device=device, dtype=torch.float32) + std_torch = torch.sqrt(torch.tensor(self.var + self.epsilon, device=device, dtype=torch.float32)) + return torch.clamp((input - mean_torch) / std_torch, -self.clip_obs, self.clip_obs) + + def update_normalizer(self, rollouts, expert_loader): + policy_data_generator = rollouts.feed_forward_generator_amp(None, mini_batch_size=expert_loader.batch_size) + expert_data_generator = expert_loader.dataset.feed_forward_generator_amp(expert_loader.batch_size) + + for expert_batch, policy_batch in zip(expert_data_generator, policy_data_generator): + self.update(torch.vstack(tuple(policy_batch) + tuple(expert_batch)).cpu().numpy()) + + +def resolve_nn_activation(act_name: str) -> torch.nn.Module: + """Resolves the activation function from the name. + + Args: + act_name: The name of the activation function. + + Returns: + The activation function. + + Raises: + ValueError: If the activation function is not found. + """ + act_dict = { + "elu": torch.nn.ELU(), + "selu": torch.nn.SELU(), + "relu": torch.nn.ReLU(), + "crelu": torch.nn.CELU(), + "lrelu": torch.nn.LeakyReLU(), + "tanh": torch.nn.Tanh(), + "sigmoid": torch.nn.Sigmoid(), + "softplus": torch.nn.Softplus(), + "gelu": torch.nn.GELU(), + "swish": torch.nn.SiLU(), + "mish": torch.nn.Mish(), + "identity": torch.nn.Identity(), + } + + act_name = act_name.lower() + if act_name in act_dict: + return act_dict[act_name] + else: + raise ValueError(f"Invalid activation function '{act_name}'. Valid activations are: {list(act_dict.keys())}") + + +def resolve_optimizer(optimizer_name: str) -> torch.optim.Optimizer: + """Resolves the optimizer from the name. + + Args: + optimizer_name: The name of the optimizer. + + Returns: + The optimizer. + + Raises: + ValueError: If the optimizer is not found. + """ + optimizer_dict = { + "adam": torch.optim.Adam, + "adamw": torch.optim.AdamW, + "sgd": torch.optim.SGD, + "rmsprop": torch.optim.RMSprop, + } + + optimizer_name = optimizer_name.lower() + if optimizer_name in optimizer_dict: + return optimizer_dict[optimizer_name] + else: + raise ValueError(f"Invalid optimizer '{optimizer_name}'. Valid optimizers are: {list(optimizer_dict.keys())}") + + +def split_and_pad_trajectories( + tensor: torch.Tensor | TensorDict, dones: torch.Tensor +) -> tuple[torch.Tensor | TensorDict, torch.Tensor]: + """Splits trajectories at done indices. Then concatenates them and pads with zeros up to the length of the longest + trajectory. Returns masks corresponding to valid parts of the trajectories. + + Example: + Input: [[a1, a2, a3, a4 | a5, a6], + [b1, b2 | b3, b4, b5 | b6]] + + Output:[[a1, a2, a3, a4], | [[True, True, True, True], + [a5, a6, 0, 0], | [True, True, False, False], + [b1, b2, 0, 0], | [True, True, False, False], + [b3, b4, b5, 0], | [True, True, True, False], + [b6, 0, 0, 0]] | [True, False, False, False]] + + Assumes that the input has the following order of dimensions: [time, number of envs, additional dimensions] + """ + + dones = dones.clone() + dones[-1] = 1 + # Permute the buffers to have order (num_envs, num_transitions_per_env, ...), for correct reshaping + flat_dones = dones.transpose(1, 0).reshape(-1, 1) + # Get length of trajectory by counting the number of successive not done elements + done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero()[:, 0])) + trajectory_lengths = done_indices[1:] - done_indices[:-1] + trajectory_lengths_list = trajectory_lengths.tolist() + # Extract the individual trajectories + if isinstance(tensor, TensorDict): + padded_trajectories = {} + for k, v in tensor.items(): + # split the tensor into trajectories + trajectories = torch.split(v.transpose(1, 0).flatten(0, 1), trajectory_lengths_list) + # add at least one full length trajectory + trajectories = trajectories + (torch.zeros(v.shape[0], *v.shape[2:], device=v.device),) + # pad the trajectories to the length of the longest trajectory + padded_trajectories[k] = torch.nn.utils.rnn.pad_sequence(trajectories) + # remove the added tensor + padded_trajectories[k] = padded_trajectories[k][:, :-1] + padded_trajectories = TensorDict( + padded_trajectories, batch_size=[tensor.batch_size[0], len(trajectory_lengths_list)] + ) + else: + # split the tensor into trajectories + trajectories = torch.split(tensor.transpose(1, 0).flatten(0, 1), trajectory_lengths_list) + # add at least one full length trajectory + trajectories = trajectories + (torch.zeros(tensor.shape[0], *tensor.shape[2:], device=tensor.device),) + # pad the trajectories to the length of the longest trajectory + padded_trajectories = torch.nn.utils.rnn.pad_sequence(trajectories) + # remove the added tensor + padded_trajectories = padded_trajectories[:, :-1] + # create masks for the valid parts of the trajectories + trajectory_masks = trajectory_lengths > torch.arange(0, tensor.shape[0], device=tensor.device).unsqueeze(1) + return padded_trajectories, trajectory_masks + + +def unpad_trajectories(trajectories, masks): + """Does the inverse operation of split_and_pad_trajectories()""" + # Need to transpose before and after the masking to have proper reshaping + return ( + trajectories.transpose(1, 0)[masks.transpose(1, 0)] + .view(-1, trajectories.shape[0], trajectories.shape[-1]) + .transpose(1, 0) + ) + + +def store_code_state(logdir, repositories) -> list: + git_log_dir = os.path.join(logdir, "git") + os.makedirs(git_log_dir, exist_ok=True) + file_paths = [] + for repository_file_path in repositories: + try: + repo = git.Repo(repository_file_path, search_parent_directories=True) + t = repo.head.commit.tree + except Exception: + print(f"Could not find git repository in {repository_file_path}. Skipping.") + # skip if not a git repository + continue + # get the name of the repository + repo_name = pathlib.Path(repo.working_dir).name + diff_file_name = os.path.join(git_log_dir, f"{repo_name}.diff") + # check if the diff file already exists + if os.path.isfile(diff_file_name): + continue + # write the diff file + print(f"Storing git diff for '{repo_name}' in: {diff_file_name}") + with open(diff_file_name, "x", encoding="utf-8") as f: + content = f"--- git status ---\n{repo.git.status()} \n\n\n--- git diff ---\n{repo.git.diff(t)}" + f.write(content) + # add the file path to the list of files to be uploaded + file_paths.append(diff_file_name) + return file_paths + + +def string_to_callable(name: str) -> Callable: + """Resolves the module and function names to return the function. + + Args: + name: The function name. The format should be 'module:attribute_name'. + + Raises: + ValueError: When the resolved attribute is not a function. + ValueError: When unable to resolve the attribute. + + Returns: + The function loaded from the module. + """ + try: + mod_name, attr_name = name.split(":") + mod = importlib.import_module(mod_name) + callable_object = getattr(mod, attr_name) + # check if attribute is callable + if callable(callable_object): + return callable_object + else: + raise ValueError(f"The imported object is not callable: '{name}'") + except AttributeError as e: + msg = ( + "We could not interpret the entry as a callable object. The format of input should be" + f" 'module:attribute_name'\nWhile processing input '{name}', received the error:\n {e}." + ) + raise ValueError(msg) + + +def resolve_obs_groups( + obs: TensorDict, obs_groups: dict[str, list[str]], default_sets: list[str] +) -> dict[str, list[str]]: + """Validates the observation configuration and defaults missing observation sets. + + The input is an observation dictionary `obs` containing observation groups and a configuration dictionary + `obs_groups` where the keys are the observation sets and the values are lists of observation groups. + + The configuration dictionary could for example look like: + { + "policy": ["group_1", "group_2"], + "critic": ["group_1", "group_3"] + } + + This means that the 'policy' observation set will contain the observations "group_1" and "group_2" and the + 'critic' observation set will contain the observations "group_1" and "group_3". This function will check that all + the observations in the 'policy' and 'critic' observation sets are present in the observation dictionary from the + environment. + + Additionally, if one of the `default_sets`, e.g. "critic", is not present in the configuration dictionary, + this function will: + + 1. Check if a group with the same name exists in the observations and assign this group to the observation set. + 2. If 1. fails, it will assign the observations from the 'policy' observation set to the default observation set. + + Args: + obs: Observations from the environment in the form of a dictionary. + obs_groups: Observation sets configuration. + default_sets: Reserved observation set names used by the algorithm (besides 'policy'). + If not provided in 'obs_groups', a default behavior gets triggered. + + Returns: + The resolved observation groups. + + Raises: + ValueError: If any observation set is an empty list. + ValueError: If any observation set contains an observation term that is not present in the observations. + """ + # check if policy observation set exists + if "policy" not in obs_groups.keys(): + if "policy" in obs: + obs_groups["policy"] = ["policy"] + warnings.warn( + "The observation configuration dictionary 'obs_groups' must contain the 'policy' key." + " As an observation group with the name 'policy' was found, this is assumed to be the observation set." + " Consider adding the 'policy' key to the 'obs_groups' dictionary for clarity." + " This behavior will be removed in a future version." + ) + else: + raise ValueError( + "The observation configuration dictionary 'obs_groups' must contain the 'policy' key." + f" Found keys: {list(obs_groups.keys())}" + ) + + # check all observation sets for valid observation groups + for set_name, groups in obs_groups.items(): + # check if the list is empty + if len(groups) == 0: + msg = f"The '{set_name}' key in the 'obs_groups' dictionary can not be an empty list." + if set_name in default_sets: + if set_name not in obs: + msg += " Consider removing the key to default to the observations used for the 'policy' set." + else: + msg += ( + f" Consider removing the key to default to the observation '{set_name}' from the environment." + ) + raise ValueError(msg) + # check groups exist inside the observations from the environment + for group in groups: + if group not in obs: + raise ValueError( + f"Observation '{group}' in observation set '{set_name}' not found in the observations from the" + f" environment. Available observations from the environment: {list(obs.keys())}" + ) + + # fill missing observation sets + for default_set_name in default_sets: + if default_set_name not in obs_groups.keys(): + if default_set_name in obs: + obs_groups[default_set_name] = [default_set_name] + warnings.warn( + f"The observation configuration dictionary 'obs_groups' must contain the '{default_set_name}' key." + f" As an observation group with the name '{default_set_name}' was found, this is assumed to be the" + f" observation set. Consider adding the '{default_set_name}' key to the 'obs_groups' dictionary for" + " clarity. This behavior will be removed in a future version." + ) + else: + obs_groups[default_set_name] = obs_groups["policy"].copy() + warnings.warn( + f"The observation configuration dictionary 'obs_groups' must contain the '{default_set_name}' key." + f" As the configuration for '{default_set_name}' is missing, the observations from the 'policy' set" + f" are used. Consider adding the '{default_set_name}' key to the 'obs_groups' dictionary for" + " clarity. This behavior will be removed in a future version." + ) + + # print the final parsed observation sets + print("-" * 80) + print("Resolved observation sets: ") + for set_name, groups in obs_groups.items(): + print("\t", set_name, ": ", groups) + print("-" * 80) + + return obs_groups diff --git a/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py new file mode 100644 index 0000000..243e82d --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py @@ -0,0 +1,87 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +from dataclasses import asdict +from torch.utils.tensorboard import SummaryWriter + +try: + import wandb +except ModuleNotFoundError: + raise ModuleNotFoundError("Wandb is required to log to Weights and Biases.") + + +class WandbSummaryWriter(SummaryWriter): + """Summary writer for Weights and Biases.""" + + def __init__(self, log_dir: str, flush_secs: int, cfg): + super().__init__(log_dir, flush_secs) + + # Get the run name + run_name = os.path.split(log_dir)[-1] + + try: + project = cfg["wandb_project"] + except KeyError: + raise KeyError("Please specify wandb_project in the runner config, e.g. legged_gym.") + + try: + entity = os.environ["WANDB_USERNAME"] + except KeyError: + entity = None + + # Initialize wandb + wandb.init(project=project, entity=entity, name=run_name) + + # Add log directory to wandb + wandb.config.update({"log_dir": log_dir}) + + self.name_map = { + "Train/mean_reward/time": "Train/mean_reward_time", + "Train/mean_episode_length/time": "Train/mean_episode_length_time", + } + + def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): + wandb.config.update({"runner_cfg": runner_cfg}) + wandb.config.update({"policy_cfg": policy_cfg}) + wandb.config.update({"alg_cfg": alg_cfg}) + try: + wandb.config.update({"env_cfg": env_cfg.to_dict()}) + except Exception: + wandb.config.update({"env_cfg": asdict(env_cfg)}) + + def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_style=False): + super().add_scalar( + tag, + scalar_value, + global_step=global_step, + walltime=walltime, + new_style=new_style, + ) + wandb.log({self._map_path(tag): scalar_value}, step=global_step) + + def stop(self): + wandb.finish() + + def log_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): + self.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg) + + def save_model(self, model_path, iter): + wandb.save(model_path, base_path=os.path.dirname(model_path)) + + def save_file(self, path, iter=None): + wandb.save(path, base_path=os.path.dirname(path)) + + """ + Private methods. + """ + + def _map_path(self, path): + if path in self.name_map: + return self.name_map[path] + else: + return path diff --git a/rsl_rl/build/lib/rsl_rl/env/__init__.py b/rsl_rl/build/lib/rsl_rl/env/__init__.py new file mode 100644 index 0000000..ab7c056 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/env/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Submodule defining the environment definitions.""" + +from .vec_env import VecEnv + +__all__ = ["VecEnv"] diff --git a/rsl_rl/build/lib/rsl_rl/env/vec_env.py b/rsl_rl/build/lib/rsl_rl/env/vec_env.py new file mode 100644 index 0000000..ae471cf --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/env/vec_env.py @@ -0,0 +1,113 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from abc import ABC, abstractmethod +from tensordict import TensorDict + + +class VecEnv(ABC): + """Abstract class for a vectorized environment. + + The vectorized environment is a collection of environments that are synchronized. This means that + the same type of action is applied to all environments and the same type of observation is returned from all + environments. + """ + + num_envs: int + """Number of environments.""" + + num_actions: int + """Number of actions.""" + + max_episode_length: int | torch.Tensor + + max_episode_length_s: float + """Maximum episode length. + + The maximum episode length can be a scalar or a tensor. If it is a scalar, it is the same for all environments. + If it is a tensor, it is the maximum episode length for each environment. This is useful for dynamic episode + lengths. + """ + + episode_length_buf: torch.Tensor + """Buffer for current episode lengths.""" + + device: torch.device | str + """Device to use.""" + + cfg: dict | object + """Configuration object.""" + + reset_env_ids: torch.Tensor | None = None + + contact_phase: torch.Tensor | None = None + """ + Operations. + """ + + @abstractmethod + def get_observations(self) -> TensorDict: + """Return the current observations. + + Returns: + observations (TensorDict): Observations from the environment. + """ + raise NotImplementedError + + @abstractmethod + def get_amp_observations(self) -> TensorDict: + """Return the current AMP observations. + + Returns: + observations (TensorDict): Observations from the environment. + """ + raise NotImplementedError + + @abstractmethod + def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]: + """Apply input action to the environment. + + Args: + actions (torch.Tensor): Input actions to apply. Shape: (num_envs, num_actions) + + Returns: + observations (TensorDict): Observations from the environment. + rewards (torch.Tensor): Rewards from the environment. Shape: (num_envs,) + dones (torch.Tensor): Done flags from the environment. Shape: (num_envs,) + extras (dict): Extra information from the environment. + + Observations: + + The observations TensorDict usually contains multiple observation groups. The `obs_groups` + dictionary of the runner configuration specifies which observation groups are used for which + purpose, i.e., it maps the available observation groups to observation sets. The observation sets + (keys of the `obs_groups` dictionary) currently used by rsl_rl are: + + - "policy": Specified observation groups are used as input to the actor/student network. + - "critic": Specified observation groups are used as input to the critic network. + - "teacher": Specified observation groups are used as input to the teacher network. + - "rnd_state": Specified observation groups are used as input to the RND network. + + Incomplete or incorrect configurations are handled in the `resolve_obs_groups()` function in + `rsl_rl/utils/utils.py`. + + Extras: + + The extras dictionary includes metrics such as the episode reward, episode length, etc. The following + dictionary keys are used by rsl_rl: + + - "time_outs" (torch.Tensor): Timeouts for the environments. These correspond to terminations that + happen due to time limits and not due to the environment reaching a terminal state. This is useful + for environments that have a fixed episode length. + + - "log" (dict[str, float | torch.Tensor]): Additional information for logging and debugging purposes. + The key should be a string and start with "/" for namespacing. The value can be a scalar or a + tensor. If it is a tensor, the mean of the tensor is used for logging. + """ + raise NotImplementedError + \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/modules/__init__.py b/rsl_rl/build/lib/rsl_rl/modules/__init__.py new file mode 100644 index 0000000..9afe1a2 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/modules/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Definitions for neural-network components for RL-agents.""" + +from .actor_critic import ActorCritic +from .actor_critic_recurrent import ActorCriticRecurrent +from .rnd import * +from .student_teacher import StudentTeacher +from .student_teacher_recurrent import StudentTeacherRecurrent +from .symmetry import * +from .discriminator_multi import DiscriminatorMulti +__all__ = [ + "ActorCritic", + "ActorCriticRecurrent", + "StudentTeacher", + "StudentTeacherRecurrent", + "DiscriminatorMulti", +] diff --git a/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py b/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py new file mode 100644 index 0000000..0efc36b --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py @@ -0,0 +1,195 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +from torch.distributions import Normal + +from rsl_rl.networks import MLP, EmpiricalNormalization + + +class ActorCritic(nn.Module): + is_recurrent = False + + def __init__( + self, + obs, + obs_groups, + num_actions, + actor_obs_normalization=False, + critic_obs_normalization=False, + actor_hidden_dims=[256, 256, 256], + critic_hidden_dims=[256, 256, 256], + activation="elu", + init_noise_std=1.0, + noise_std_type: str = "scalar", + state_dependent_std=False, + **kwargs, + ): + if kwargs: + print( + "ActorCritic.__init__ got unexpected arguments, which will be ignored: " + + str([key for key in kwargs.keys()]) + ) + super().__init__() + + # get the observation dimensions + self.obs_groups = obs_groups + num_actor_obs = 0 + for obs_group in obs_groups["policy"]: + assert len(obs[obs_group].shape) == 2, "The ActorCritic module only supports 1D observations." + num_actor_obs += obs[obs_group].shape[-1] + num_critic_obs = 0 + for obs_group in obs_groups["critic"]: + assert len(obs[obs_group].shape) == 2, "The ActorCritic module only supports 1D observations." + num_critic_obs += obs[obs_group].shape[-1] + + self.state_dependent_std = state_dependent_std + # actor + if self.state_dependent_std: + self.actor = MLP(num_actor_obs, [2, num_actions], actor_hidden_dims, activation) + else: + self.actor = MLP(num_actor_obs, num_actions, actor_hidden_dims, activation) + # actor observation normalization + self.actor_obs_normalization = actor_obs_normalization + if actor_obs_normalization: + self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs) + else: + self.actor_obs_normalizer = torch.nn.Identity() + print(f"Actor MLP: {self.actor}") + + # critic + self.critic = MLP(num_critic_obs, 1, critic_hidden_dims, activation) + # critic observation normalization + self.critic_obs_normalization = critic_obs_normalization + if critic_obs_normalization: + self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs) + else: + self.critic_obs_normalizer = torch.nn.Identity() + print(f"Critic MLP: {self.critic}") + + # Action noise + self.noise_std_type = noise_std_type + if self.state_dependent_std: + torch.nn.init.zeros_(self.actor[-2].weight[num_actions:]) + if self.noise_std_type == "scalar": + torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std) + elif self.noise_std_type == "log": + torch.nn.init.constant_( + self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7)) + ) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + else: + if self.noise_std_type == "scalar": + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + elif self.noise_std_type == "log": + self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + + # Action distribution (populated in update_distribution) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args(False) + + def reset(self, dones=None): + pass + + def forward(self): + raise NotImplementedError + + @property + def action_mean(self): + return self.distribution.mean + + @property + def action_std(self): + return self.distribution.stddev + + @property + def entropy(self): + return self.distribution.entropy().sum(dim=-1) + + def update_distribution(self, obs): + if self.state_dependent_std: + # compute mean and standard deviation + mean_and_std = self.actor(obs) + if self.noise_std_type == "scalar": + mean, std = torch.unbind(mean_and_std, dim=-2) + elif self.noise_std_type == "log": + mean, log_std = torch.unbind(mean_and_std, dim=-2) + std = torch.exp(log_std) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + else: + # compute mean + mean = self.actor(obs) + # compute standard deviation + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + # create distribution + self.distribution = Normal(mean, std) + + def act(self, obs, **kwargs): + obs = self.get_actor_obs(obs) + obs = self.actor_obs_normalizer(obs) + self.update_distribution(obs) + return self.distribution.sample() + + def act_inference(self, obs): + obs = self.get_actor_obs(obs) + obs = self.actor_obs_normalizer(obs) + return self.actor(obs) + + def evaluate(self, obs, **kwargs): + obs = self.get_critic_obs(obs) + obs = self.critic_obs_normalizer(obs) + return self.critic(obs) + + def get_actor_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["policy"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_critic_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["critic"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_actions_log_prob(self, actions): + return self.distribution.log_prob(actions).sum(dim=-1) + + def update_normalization(self, obs): + if self.actor_obs_normalization: + actor_obs = self.get_actor_obs(obs) + self.actor_obs_normalizer.update(actor_obs) + if self.critic_obs_normalization: + critic_obs = self.get_critic_obs(obs) + self.critic_obs_normalizer.update(critic_obs) + + def load_state_dict(self, state_dict, strict=True): + """Load the parameters of the actor-critic model. + + Args: + state_dict (dict): State dictionary of the model. + strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this + module's state_dict() function. + + Returns: + bool: Whether this training resumes a previous training. This flag is used by the `load()` function of + `OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation). + """ + + super().load_state_dict(state_dict, strict=strict) + return True # training resumes diff --git a/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py b/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py new file mode 100644 index 0000000..bba46ca --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py @@ -0,0 +1,218 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +import warnings +from torch.distributions import Normal + +from rsl_rl.networks import MLP, EmpiricalNormalization, Memory + + +class ActorCriticRecurrent(nn.Module): + is_recurrent = True + + def __init__( + self, + obs, + obs_groups, + num_actions, + actor_obs_normalization=False, + critic_obs_normalization=False, + actor_hidden_dims=[256, 256, 256], + critic_hidden_dims=[256, 256, 256], + activation="elu", + init_noise_std=1.0, + noise_std_type: str = "scalar", + state_dependent_std=False, + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=1, + **kwargs, + ): + if "rnn_hidden_size" in kwargs: + warnings.warn( + "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. " + "Please use `rnn_hidden_dim` instead.", + DeprecationWarning, + ) + if rnn_hidden_dim == 256: # Only override if the new argument is at its default + rnn_hidden_dim = kwargs.pop("rnn_hidden_size") + if kwargs: + print( + "ActorCriticRecurrent.__init__ got unexpected arguments, which will be ignored: " + str(kwargs.keys()), + ) + super().__init__() + + # get the observation dimensions + self.obs_groups = obs_groups + num_actor_obs = 0 + for obs_group in obs_groups["policy"]: + assert len(obs[obs_group].shape) == 2, "The ActorCriticRecurrent module only supports 1D observations." + num_actor_obs += obs[obs_group].shape[-1] + num_critic_obs = 0 + for obs_group in obs_groups["critic"]: + assert len(obs[obs_group].shape) == 2, "The ActorCriticRecurrent module only supports 1D observations." + num_critic_obs += obs[obs_group].shape[-1] + + self.state_dependent_std = state_dependent_std + # actor + self.memory_a = Memory(num_actor_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) + if self.state_dependent_std: + self.actor = MLP(rnn_hidden_dim, [2, num_actions], actor_hidden_dims, activation) + else: + self.actor = MLP(rnn_hidden_dim, num_actions, actor_hidden_dims, activation) + + # actor observation normalization + self.actor_obs_normalization = actor_obs_normalization + if actor_obs_normalization: + self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs) + else: + self.actor_obs_normalizer = torch.nn.Identity() + print(f"Actor RNN: {self.memory_a}") + print(f"Actor MLP: {self.actor}") + + # critic + self.memory_c = Memory(num_critic_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) + self.critic = MLP(rnn_hidden_dim, 1, critic_hidden_dims, activation) + # critic observation normalization + self.critic_obs_normalization = critic_obs_normalization + if critic_obs_normalization: + self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs) + else: + self.critic_obs_normalizer = torch.nn.Identity() + print(f"Critic RNN: {self.memory_c}") + print(f"Critic MLP: {self.critic}") + + # Action noise + self.noise_std_type = noise_std_type + if self.state_dependent_std: + torch.nn.init.zeros_(self.actor[-2].weight[num_actions:]) + if self.noise_std_type == "scalar": + torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std) + elif self.noise_std_type == "log": + torch.nn.init.constant_( + self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7)) + ) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + else: + if self.noise_std_type == "scalar": + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + elif self.noise_std_type == "log": + self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + + # Action distribution (populated in update_distribution) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args(False) + + @property + def action_mean(self): + return self.distribution.mean + + @property + def action_std(self): + return self.distribution.stddev + + @property + def entropy(self): + return self.distribution.entropy().sum(dim=-1) + + def reset(self, dones=None): + self.memory_a.reset(dones) + self.memory_c.reset(dones) + + def forward(self): + raise NotImplementedError + + def update_distribution(self, obs): + if self.state_dependent_std: + # compute mean and standard deviation + mean_and_std = self.actor(obs) + if self.noise_std_type == "scalar": + mean, std = torch.unbind(mean_and_std, dim=-2) + elif self.noise_std_type == "log": + mean, log_std = torch.unbind(mean_and_std, dim=-2) + std = torch.exp(log_std) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + else: + # compute mean + mean = self.actor(obs) + # compute standard deviation + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + # create distribution + self.distribution = Normal(mean, std) + + def act(self, obs, masks=None, hidden_states=None): + obs = self.get_actor_obs(obs) + obs = self.actor_obs_normalizer(obs) + out_mem = self.memory_a(obs, masks, hidden_states).squeeze(0) + self.update_distribution(out_mem) + return self.distribution.sample() + + def act_inference(self, obs): + obs = self.get_actor_obs(obs) + obs = self.actor_obs_normalizer(obs) + out_mem = self.memory_a(obs).squeeze(0) + return self.actor(out_mem) + + def evaluate(self, obs, masks=None, hidden_states=None): + obs = self.get_critic_obs(obs) + obs = self.critic_obs_normalizer(obs) + out_mem = self.memory_c(obs, masks, hidden_states).squeeze(0) + return self.critic(out_mem) + + def get_actor_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["policy"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_critic_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["critic"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_actions_log_prob(self, actions): + return self.distribution.log_prob(actions).sum(dim=-1) + + def get_hidden_states(self): + return self.memory_a.hidden_states, self.memory_c.hidden_states + + def update_normalization(self, obs): + if self.actor_obs_normalization: + actor_obs = self.get_actor_obs(obs) + self.actor_obs_normalizer.update(actor_obs) + if self.critic_obs_normalization: + critic_obs = self.get_critic_obs(obs) + self.critic_obs_normalizer.update(critic_obs) + + def load_state_dict(self, state_dict, strict=True): + """Load the parameters of the actor-critic model. + + Args: + state_dict (dict): State dictionary of the model. + strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this + module's state_dict() function. + + Returns: + bool: Whether this training resumes a previous training. This flag is used by the `load()` function of + `OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation). + """ + + super().load_state_dict(state_dict, strict=strict) + return True diff --git a/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py b/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py new file mode 100644 index 0000000..7899d23 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py @@ -0,0 +1,102 @@ +import torch +import torch.nn as nn +from torch import autograd +import torch.nn.utils.spectral_norm as spectral_norm + + +class DiscriminatorMulti(nn.Module): + def __init__( + self, state_dim, amp_reward_coef, hidden_layer_sizes, device, + num_frames=2, task_reward_lerp=0.0, use_lerp=True): + super(DiscriminatorMulti, self).__init__() + + self.device = device + self.state_dim = state_dim + self.use_lerp = use_lerp + self.num_frames = num_frames # 存储帧数参数 + + self.amp_reward_coef = amp_reward_coef + amp_layers = [] + + curr_in_dim = state_dim * num_frames + for hidden_dim in hidden_layer_sizes: + amp_layers.append(spectral_norm(nn.Linear(curr_in_dim, hidden_dim))) + amp_layers.append(nn.ReLU()) + curr_in_dim = hidden_dim + self.trunk = nn.Sequential(*amp_layers).to(device) + self.amp_linear = spectral_norm(nn.Linear(hidden_layer_sizes[-1], 1)).to(device) + + self.trunk.train() + self.amp_linear.train() + + self.task_reward_lerp = task_reward_lerp + + def forward(self, x): + h = self.trunk(x) + d = self.amp_linear(h) + return d + + def compute_grad_pen(self, + expert_states, # 改为接收多帧状态列表 + lambda_=10): + # 将多帧状态沿最后一个维度拼接 + expert_data = expert_states.flatten(1) + expert_data.requires_grad = True + + disc = self.amp_linear(self.trunk(expert_data)) + ones = torch.ones(disc.size(), device=disc.device) + grad = autograd.grad( + outputs=disc, inputs=expert_data, + grad_outputs=ones, create_graph=True, + retain_graph=True, only_inputs=True)[0] + + # Enforce that the grad norm approaches 0. + grad_pen = lambda_ * (grad.norm(2, dim=1) - 0).pow(2).mean() + return grad_pen + + + def get_disc_weights(self): + weights = [] + for m in self.trunk.modules(): + if isinstance(m, nn.Linear): + weights.append(torch.flatten(m.weight)) + + weights.append(torch.flatten(self.amp_linear.weight)) + return weights + + def get_disc_logit_weights(self): + return torch.flatten(self.amp_linear.weight) + + def predict_amp_reward( + self, states, # 改为接收多帧状态列表 + task_reward, normalizer=None): + """ + states: torch.Tensor, shape=(num_envs, num_frames, state_dim) + task_reward: torch.Tensor, shape=(num_envs, 1) + """ + # import ipdb; ipdb.set_trace() + with torch.no_grad(): + self.eval() + if normalizer is not None: + # 对每一帧状态进行归一化 + states = normalizer.normalize_torch(states, self.device) + + # 拼接多帧状态 + state_cat = states.flatten(1) + d = self.amp_linear(self.trunk(state_cat)) + disc_reward = self.amp_reward_coef * torch.clamp(1 - (1/4) * torch.square(d - 1), min=0) + + if self.use_lerp: + if self.task_reward_lerp > 0: + reward = self._lerp_reward(disc_reward, task_reward.unsqueeze(-1)) + self.train() + return reward.squeeze(), d, disc_reward.squeeze() * (1.0 - self.task_reward_lerp) + else: + disc_reward *= 0.02 + reward = task_reward.unsqueeze(-1) + disc_reward + self.train() + return reward.squeeze(), d, disc_reward.squeeze() + + def _lerp_reward(self, disc_r, task_r): + r = (1.0 - self.task_reward_lerp) * disc_r + self.task_reward_lerp * task_r + return r \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/modules/rnd.py b/rsl_rl/build/lib/rsl_rl/modules/rnd.py new file mode 100644 index 0000000..8e65c43 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/modules/rnd.py @@ -0,0 +1,209 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn + +from rsl_rl.networks import MLP, EmpiricalDiscountedVariationNormalization, EmpiricalNormalization + + +class RandomNetworkDistillation(nn.Module): + """Implementation of Random Network Distillation (RND) [1] + + References: + .. [1] Burda, Yuri, et al. "Exploration by random network distillation." arXiv preprint arXiv:1810.12894 (2018). + """ + + def __init__( + self, + num_states: int, + obs_groups: dict, + num_outputs: int, + predictor_hidden_dims: list[int], + target_hidden_dims: list[int], + activation: str = "elu", + weight: float = 0.0, + state_normalization: bool = False, + reward_normalization: bool = False, + device: str = "cpu", + weight_schedule: dict | None = None, + ): + """Initialize the RND module. + + - If :attr:`state_normalization` is True, then the input state is normalized using an Empirical Normalization layer. + - If :attr:`reward_normalization` is True, then the intrinsic reward is normalized using an Empirical Discounted + Variation Normalization layer. + + .. note:: + If the hidden dimensions are -1 in the predictor and target networks configuration, then the number of states + is used as the hidden dimension. + + Args: + num_states: Number of states/inputs to the predictor and target networks. + num_outputs: Number of outputs (embedding size) of the predictor and target networks. + predictor_hidden_dims: List of hidden dimensions of the predictor network. + target_hidden_dims: List of hidden dimensions of the target network. + activation: Activation function. Defaults to "elu". + weight: Scaling factor of the intrinsic reward. Defaults to 0.0. + state_normalization: Whether to normalize the input state. Defaults to False. + reward_normalization: Whether to normalize the intrinsic reward. Defaults to False. + device: Device to use. Defaults to "cpu". + weight_schedule: The type of schedule to use for the RND weight parameter. + Defaults to None, in which case the weight parameter is constant. + It is a dictionary with the following keys: + + - "mode": The type of schedule to use for the RND weight parameter. + - "constant": Constant weight schedule. + - "step": Step weight schedule. + - "linear": Linear weight schedule. + + For the "step" weight schedule, the following parameters are required: + + - "final_step": The step at which the weight parameter is set to the final value. + - "final_value": The final value of the weight parameter. + + For the "linear" weight schedule, the following parameters are required: + - "initial_step": The step at which the weight parameter is set to the initial value. + - "final_step": The step at which the weight parameter is set to the final value. + - "final_value": The final value of the weight parameter. + """ + # initialize parent class + super().__init__() + + # Store parameters + self.num_states = num_states + self.obs_groups = obs_groups + self.num_outputs = num_outputs + self.initial_weight = weight + self.device = device + self.state_normalization = state_normalization + self.reward_normalization = reward_normalization + + # Normalization of input gates + if state_normalization: + self.state_normalizer = EmpiricalNormalization(shape=[self.num_states], until=1.0e8).to(self.device) + else: + self.state_normalizer = torch.nn.Identity() + # Normalization of intrinsic reward + if reward_normalization: + self.reward_normalizer = EmpiricalDiscountedVariationNormalization(shape=[], until=1.0e8).to(self.device) + else: + self.reward_normalizer = torch.nn.Identity() + + # counter for the number of updates + self.update_counter = 0 + + # resolve weight schedule + if weight_schedule is not None: + self.weight_scheduler_params = weight_schedule + self.weight_scheduler = getattr(self, f"_{weight_schedule['mode']}_weight_schedule") + else: + self.weight_scheduler = None + # Create network architecture + self.predictor = MLP(num_states, num_outputs, predictor_hidden_dims, activation).to(self.device) + self.target = MLP(num_states, num_outputs, target_hidden_dims, activation).to(self.device) + + # make target network not trainable + self.target.eval() + + def get_intrinsic_reward(self, obs) -> torch.Tensor: + # Note: the counter is updated number of env steps per learning iteration + self.update_counter += 1 + # Extract the rnd state from the observation + rnd_state = self.get_rnd_state(obs) + rnd_state = self.state_normalizer(rnd_state) + # Obtain the embedding of the rnd state from the target and predictor networks + target_embedding = self.target(rnd_state).detach() + predictor_embedding = self.predictor(rnd_state).detach() + # Compute the intrinsic reward as the distance between the embeddings + intrinsic_reward = torch.linalg.norm(target_embedding - predictor_embedding, dim=1) + # Normalize intrinsic reward + intrinsic_reward = self.reward_normalizer(intrinsic_reward) + + # Check the weight schedule + if self.weight_scheduler is not None: + self.weight = self.weight_scheduler(step=self.update_counter, **self.weight_scheduler_params) + else: + self.weight = self.initial_weight + # Scale intrinsic reward + intrinsic_reward *= self.weight + + return intrinsic_reward + + def forward(self, *args, **kwargs): + raise RuntimeError("Forward method is not implemented. Use get_intrinsic_reward instead.") + + def train(self, mode: bool = True): + # sets module into training mode + self.predictor.train(mode) + if self.state_normalization: + self.state_normalizer.train(mode) + if self.reward_normalization: + self.reward_normalizer.train(mode) + return self + + def eval(self): + return self.train(False) + + def get_rnd_state(self, obs): + obs_list = [] + for obs_group in self.obs_groups["rnd_state"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def update_normalization(self, obs): + # Normalize the state + if self.state_normalization: + rnd_state = self.get_rnd_state(obs) + self.state_normalizer.update(rnd_state) + + """ + Different weight schedules. + """ + + def _constant_weight_schedule(self, step: int, **kwargs): + return self.initial_weight + + def _step_weight_schedule(self, step: int, final_step: int, final_value: float, **kwargs): + return self.initial_weight if step < final_step else final_value + + def _linear_weight_schedule(self, step: int, initial_step: int, final_step: int, final_value: float, **kwargs): + if step < initial_step: + return self.initial_weight + elif step > final_step: + return final_value + else: + return self.initial_weight + (final_value - self.initial_weight) * (step - initial_step) / ( + final_step - initial_step + ) + + +def resolve_rnd_config(alg_cfg, obs, obs_groups, env): + """Resolve the RND configuration. + + Args: + alg_cfg: The algorithm configuration dictionary. + obs: The observation dictionary. + obs_groups: The observation groups dictionary. + env: The environment. + + Returns: + The resolved algorithm configuration dictionary. + """ + # resolve dimension of rnd gated state + if "rnd_cfg" in alg_cfg and alg_cfg["rnd_cfg"] is not None: + # get dimension of rnd gated state + num_rnd_state = 0 + for obs_group in obs_groups["rnd_state"]: + assert len(obs[obs_group].shape) == 2, "The RND module only supports 1D observations." + num_rnd_state += obs[obs_group].shape[-1] + # add rnd gated state to config + alg_cfg["rnd_cfg"]["num_states"] = num_rnd_state + alg_cfg["rnd_cfg"]["obs_groups"] = obs_groups + # scale down the rnd weight with timestep + alg_cfg["rnd_cfg"]["weight"] *= env.unwrapped.step_dt + return alg_cfg diff --git a/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py b/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py new file mode 100644 index 0000000..6bf1380 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py @@ -0,0 +1,206 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +from torch.distributions import Normal + +from rsl_rl.networks import MLP, EmpiricalNormalization + + +class StudentTeacher(nn.Module): + is_recurrent = False + + def __init__( + self, + obs, + obs_groups, + num_actions, + student_obs_normalization=False, + teacher_obs_normalization=False, + student_hidden_dims=[256, 256, 256], + teacher_hidden_dims=[256, 256, 256], + activation="elu", + init_noise_std=0.1, + noise_std_type: str = "scalar", + **kwargs, + ): + if kwargs: + print( + "StudentTeacher.__init__ got unexpected arguments, which will be ignored: " + + str([key for key in kwargs.keys()]) + ) + super().__init__() + + self.loaded_teacher = False # indicates if teacher has been loaded + + # get the observation dimensions + self.obs_groups = obs_groups + num_student_obs = 0 + for obs_group in obs_groups["policy"]: + assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." + num_student_obs += obs[obs_group].shape[-1] + num_teacher_obs = 0 + for obs_group in obs_groups["teacher"]: + assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." + num_teacher_obs += obs[obs_group].shape[-1] + + # student + self.student = MLP(num_student_obs, num_actions, student_hidden_dims, activation) + + # student observation normalization + self.student_obs_normalization = student_obs_normalization + if student_obs_normalization: + self.student_obs_normalizer = EmpiricalNormalization(num_student_obs) + else: + self.student_obs_normalizer = torch.nn.Identity() + + print(f"Student MLP: {self.student}") + + # teacher + self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation) + self.teacher.eval() + + # teacher observation normalization + self.teacher_obs_normalization = teacher_obs_normalization + if teacher_obs_normalization: + self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs) + else: + self.teacher_obs_normalizer = torch.nn.Identity() + + print(f"Teacher MLP: {self.teacher}") + + # action noise + self.noise_std_type = noise_std_type + if self.noise_std_type == "scalar": + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + elif self.noise_std_type == "log": + self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + + # action distribution (populated in update_distribution) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args(False) + + def reset(self, dones=None, hidden_states=None): + pass + + def forward(self): + raise NotImplementedError + + @property + def action_mean(self): + return self.distribution.mean + + @property + def action_std(self): + return self.distribution.stddev + + @property + def entropy(self): + return self.distribution.entropy().sum(dim=-1) + + def update_distribution(self, obs): + # compute mean + mean = self.student(obs) + # compute standard deviation + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + # create distribution + self.distribution = Normal(mean, std) + + def act(self, obs): + obs = self.get_student_obs(obs) + obs = self.student_obs_normalizer(obs) + self.update_distribution(obs) + return self.distribution.sample() + + def act_inference(self, obs): + obs = self.get_student_obs(obs) + obs = self.student_obs_normalizer(obs) + return self.student(obs) + + def evaluate(self, obs): + obs = self.get_teacher_obs(obs) + obs = self.teacher_obs_normalizer(obs) + with torch.no_grad(): + return self.teacher(obs) + + def get_student_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["policy"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_teacher_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["teacher"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_hidden_states(self): + return None + + def detach_hidden_states(self, dones=None): + pass + + def train(self, mode=True): + super().train(mode) + # make sure teacher is in eval mode + self.teacher.eval() + self.teacher_obs_normalizer.eval() + + def update_normalization(self, obs): + if self.student_obs_normalization: + student_obs = self.get_student_obs(obs) + self.student_obs_normalizer.update(student_obs) + + def load_state_dict(self, state_dict, strict=True): + """Load the parameters of the student and teacher networks. + + Args: + state_dict (dict): State dictionary of the model. + strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this + module's state_dict() function. + + Returns: + bool: Whether this training resumes a previous training. This flag is used by the `load()` function of + `OnPolicyRunner` to determine how to load further parameters. + """ + + # check if state_dict contains teacher and student or just teacher parameters + if any("actor" in key for key in state_dict.keys()): # loading parameters from rl training + # rename keys to match teacher and remove critic parameters + teacher_state_dict = {} + teacher_obs_normalizer_state_dict = {} + for key, value in state_dict.items(): + if "actor." in key: + teacher_state_dict[key.replace("actor.", "")] = value + if "actor_obs_normalizer." in key: + teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value + self.teacher.load_state_dict(teacher_state_dict, strict=strict) + self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict) + # set flag for successfully loading the parameters + self.loaded_teacher = True + self.teacher.eval() + self.teacher_obs_normalizer.eval() + return False # training does not resume + elif any("student" in key for key in state_dict.keys()): # loading parameters from distillation training + super().load_state_dict(state_dict, strict=strict) + # set flag for successfully loading the parameters + self.loaded_teacher = True + self.teacher.eval() + self.teacher_obs_normalizer.eval() + return True # training resumes + else: + raise ValueError("state_dict does not contain student or teacher parameters") diff --git a/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py b/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py new file mode 100644 index 0000000..964a2dc --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py @@ -0,0 +1,249 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +import warnings +from torch.distributions import Normal + +from rsl_rl.networks import MLP, EmpiricalNormalization, Memory + + +class StudentTeacherRecurrent(nn.Module): + is_recurrent = True + + def __init__( + self, + obs, + obs_groups, + num_actions, + student_obs_normalization=False, + teacher_obs_normalization=False, + student_hidden_dims=[256, 256, 256], + teacher_hidden_dims=[256, 256, 256], + activation="elu", + init_noise_std=0.1, + noise_std_type: str = "scalar", + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=1, + teacher_recurrent=False, + **kwargs, + ): + if "rnn_hidden_size" in kwargs: + warnings.warn( + "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. " + "Please use `rnn_hidden_dim` instead.", + DeprecationWarning, + ) + if rnn_hidden_dim == 256: # Only override if the new argument is at its default + rnn_hidden_dim = kwargs.pop("rnn_hidden_size") + if kwargs: + print( + "StudentTeacherRecurrent.__init__ got unexpected arguments, which will be ignored: " + + str(kwargs.keys()), + ) + super().__init__() + + self.loaded_teacher = False # indicates if teacher has been loaded + self.teacher_recurrent = teacher_recurrent # indicates if teacher is recurrent too + + # get the observation dimensions + self.obs_groups = obs_groups + num_student_obs = 0 + for obs_group in obs_groups["policy"]: + assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." + num_student_obs += obs[obs_group].shape[-1] + num_teacher_obs = 0 + for obs_group in obs_groups["teacher"]: + assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." + num_teacher_obs += obs[obs_group].shape[-1] + + # student + self.memory_s = Memory(num_student_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) + self.student = MLP(rnn_hidden_dim, num_actions, student_hidden_dims, activation) + + # student observation normalization + self.student_obs_normalization = student_obs_normalization + if student_obs_normalization: + self.student_obs_normalizer = EmpiricalNormalization(num_student_obs) + else: + self.student_obs_normalizer = torch.nn.Identity() + + print(f"Student RNN: {self.memory_s}") + print(f"Student MLP: {self.student}") + + # teacher + if self.teacher_recurrent: + self.memory_t = Memory( + num_teacher_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim + ) + num_teacher_obs = rnn_hidden_dim + self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation) + + # teacher observation normalization + self.teacher_obs_normalization = teacher_obs_normalization + if teacher_obs_normalization: + self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs) + else: + self.teacher_obs_normalizer = torch.nn.Identity() + + if self.teacher_recurrent: + print(f"Teacher RNN: {self.memory_t}") + print(f"Teacher MLP: {self.teacher}") + + # action noise + self.noise_std_type = noise_std_type + if self.noise_std_type == "scalar": + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + elif self.noise_std_type == "log": + self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + + # action distribution (populated in update_distribution) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args(False) + + def reset(self, dones=None, hidden_states=None): + if hidden_states is None: + hidden_states = (None, None) + self.memory_s.reset(dones, hidden_states[0]) + if self.teacher_recurrent: + self.memory_t.reset(dones, hidden_states[1]) + + def forward(self): + raise NotImplementedError + + @property + def action_mean(self): + return self.distribution.mean + + @property + def action_std(self): + return self.distribution.stddev + + @property + def entropy(self): + return self.distribution.entropy().sum(dim=-1) + + def update_distribution(self, obs): + # compute mean + mean = self.student(obs) + # compute standard deviation + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + # create distribution + self.distribution = Normal(mean, std) + + def act(self, obs): + obs = self.get_student_obs(obs) + obs = self.student_obs_normalizer(obs) + out_mem = self.memory_s(obs).squeeze(0) + self.update_distribution(out_mem) + return self.distribution.sample() + + def act_inference(self, obs): + obs = self.get_student_obs(obs) + obs = self.student_obs_normalizer(obs) + out_mem = self.memory_s(obs).squeeze(0) + return self.student(out_mem) + + def evaluate(self, obs): + obs = self.get_teacher_obs(obs) + obs = self.teacher_obs_normalizer(obs) + with torch.no_grad(): + if self.teacher_recurrent: + self.memory_t.eval() + obs = self.memory_t(obs).squeeze(0) + return self.teacher(obs) + + def get_student_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["policy"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_teacher_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["teacher"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_hidden_states(self): + if self.teacher_recurrent: + return self.memory_s.hidden_states, self.memory_t.hidden_states + else: + return self.memory_s.hidden_states, None + + def detach_hidden_states(self, dones=None): + self.memory_s.detach_hidden_states(dones) + if self.teacher_recurrent: + self.memory_t.detach_hidden_states(dones) + + def train(self, mode=True): + super().train(mode) + # make sure teacher is in eval mode + self.teacher.eval() + self.teacher_obs_normalizer.eval() + + def update_normalization(self, obs): + if self.student_obs_normalization: + student_obs = self.get_student_obs(obs) + self.student_obs_normalizer.update(student_obs) + + def load_state_dict(self, state_dict, strict=True): + """Load the parameters of the student and teacher networks. + + Args: + state_dict (dict): State dictionary of the model. + strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this + module's state_dict() function. + + Returns: + bool: Whether this training resumes a previous training. This flag is used by the `load()` function of + `OnPolicyRunner` to determine how to load further parameters. + """ + + # check if state_dict contains teacher and student or just teacher parameters + if any("actor" in key for key in state_dict.keys()): # loading parameters from rl training + # rename keys to match teacher and remove critic parameters + teacher_state_dict = {} + teacher_obs_normalizer_state_dict = {} + for key, value in state_dict.items(): + if "actor." in key: + teacher_state_dict[key.replace("actor.", "")] = value + if "actor_obs_normalizer." in key: + teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value + self.teacher.load_state_dict(teacher_state_dict, strict=strict) + self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict) + # also load recurrent memory if teacher is recurrent + if self.teacher_recurrent: + memory_t_state_dict = {} + for key, value in state_dict.items(): + if "memory_a." in key: + memory_t_state_dict[key.replace("memory_a.", "")] = value + self.memory_t.load_state_dict(memory_t_state_dict, strict=strict) + # set flag for successfully loading the parameters + self.loaded_teacher = True + self.teacher.eval() + self.teacher_obs_normalizer.eval() + return False # training does not resume + elif any("student" in key for key in state_dict.keys()): # loading parameters from distillation training + super().load_state_dict(state_dict, strict=strict) + # set flag for successfully loading the parameters + self.loaded_teacher = True + self.teacher.eval() + self.teacher_obs_normalizer.eval() + return True # training resumes + else: + raise ValueError("state_dict does not contain student or teacher parameters") diff --git a/rsl_rl/build/lib/rsl_rl/modules/symmetry.py b/rsl_rl/build/lib/rsl_rl/modules/symmetry.py new file mode 100644 index 0000000..b017515 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/modules/symmetry.py @@ -0,0 +1,24 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + + +def resolve_symmetry_config(alg_cfg, env): + """Resolve the symmetry configuration. + + Args: + alg_cfg: The algorithm configuration dictionary. + env: The environment. + + Returns: + The resolved algorithm configuration dictionary. + """ + + # if using symmetry then pass the environment config object + if "symmetry_cfg" in alg_cfg and alg_cfg["symmetry_cfg"] is not None: + # this is used by the symmetry function for handling different observation terms + alg_cfg["symmetry_cfg"]["_env"] = env + return alg_cfg diff --git a/rsl_rl/build/lib/rsl_rl/networks/__init__.py b/rsl_rl/build/lib/rsl_rl/networks/__init__.py new file mode 100644 index 0000000..c18f487 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/networks/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Definitions for components of modules.""" + +from .memory import Memory +from .mlp import MLP +from .normalization import EmpiricalDiscountedVariationNormalization, EmpiricalNormalization diff --git a/rsl_rl/build/lib/rsl_rl/networks/memory.py b/rsl_rl/build/lib/rsl_rl/networks/memory.py new file mode 100644 index 0000000..7577357 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/networks/memory.py @@ -0,0 +1,70 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch.nn as nn + +from rsl_rl.utils import unpad_trajectories + + +class Memory(nn.Module): + """Memory module for recurrent networks. + + This module is used to store the hidden states of the policy. + Currently only supports GRU and LSTM. + """ + + def __init__(self, input_size, type="lstm", num_layers=1, hidden_size=256): + super().__init__() + # RNN + rnn_cls = nn.GRU if type.lower() == "gru" else nn.LSTM + self.rnn = rnn_cls(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers) + self.hidden_states = None + + def forward(self, input, masks=None, hidden_states=None): + batch_mode = masks is not None + if batch_mode: + # batch mode: needs saved hidden states + if hidden_states is None: + raise ValueError("Hidden states not passed to memory module during policy update") + out, _ = self.rnn(input, hidden_states) + out = unpad_trajectories(out, masks) + else: + # inference/distillation mode: uses hidden states of last step + out, self.hidden_states = self.rnn(input.unsqueeze(0), self.hidden_states) + return out + + def reset(self, dones=None, hidden_states=None): + if dones is None: # reset all hidden states + if hidden_states is None: + self.hidden_states = None + else: + self.hidden_states = hidden_states + elif self.hidden_states is not None: # reset hidden states of done environments + if hidden_states is None: + if isinstance(self.hidden_states, tuple): # tuple in case of LSTM + for hidden_state in self.hidden_states: + hidden_state[..., dones == 1, :] = 0.0 + else: + self.hidden_states[..., dones == 1, :] = 0.0 + else: + NotImplementedError( + "Resetting hidden states of done environments with custom hidden states is not implemented" + ) + + def detach_hidden_states(self, dones=None): + if self.hidden_states is not None: + if dones is None: # detach all hidden states + if isinstance(self.hidden_states, tuple): # tuple in case of LSTM + self.hidden_states = tuple(hidden_state.detach() for hidden_state in self.hidden_states) + else: + self.hidden_states = self.hidden_states.detach() + else: # detach hidden states of done environments + if isinstance(self.hidden_states, tuple): # tuple in case of LSTM + for hidden_state in self.hidden_states: + hidden_state[..., dones == 1, :] = hidden_state[..., dones == 1, :].detach() + else: + self.hidden_states[..., dones == 1, :] = self.hidden_states[..., dones == 1, :].detach() diff --git a/rsl_rl/build/lib/rsl_rl/networks/mlp.py b/rsl_rl/build/lib/rsl_rl/networks/mlp.py new file mode 100644 index 0000000..e91574e --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/networks/mlp.py @@ -0,0 +1,120 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +from functools import reduce + +from rsl_rl.utils import resolve_nn_activation + + +class MLP(nn.Sequential): + """Multi-layer perceptron. + + The MLP network is a sequence of linear layers and activation functions. The + last layer is a linear layer that outputs the desired dimension unless the + last activation function is specified. + + It provides additional conveniences: + + - If the hidden dimensions have a value of ``-1``, the dimension is inferred + from the input dimension. + - If the output dimension is a tuple, the output is reshaped to the desired + shape. + + """ + + def __init__( + self, + input_dim: int, + output_dim: int | tuple[int] | list[int], + hidden_dims: tuple[int] | list[int], + activation: str = "elu", + last_activation: str | None = None, + ): + """Initialize the MLP. + + Args: + input_dim: Dimension of the input. + output_dim: Dimension of the output. + hidden_dims: Dimensions of the hidden layers. A value of ``-1`` indicates + that the dimension should be inferred from the input dimension. + activation: Activation function. Defaults to "elu". + last_activation: Activation function of the last layer. Defaults to None, + in which case the last layer is linear. + """ + super().__init__() + + # resolve activation functions + activation_mod = resolve_nn_activation(activation) + last_activation_mod = resolve_nn_activation(last_activation) if last_activation is not None else None + # resolve number of hidden dims if they are -1 + hidden_dims_processed = [input_dim if dim == -1 else dim for dim in hidden_dims] + + # create layers sequentially + layers = [] + layers.append(nn.Linear(input_dim, hidden_dims_processed[0])) + layers.append(activation_mod) + + for layer_index in range(len(hidden_dims_processed) - 1): + layers.append(nn.Linear(hidden_dims_processed[layer_index], hidden_dims_processed[layer_index + 1])) + layers.append(activation_mod) + + # add last layer + if isinstance(output_dim, int): + layers.append(nn.Linear(hidden_dims_processed[-1], output_dim)) + else: + # compute the total output dimension + total_out_dim = reduce(lambda x, y: x * y, output_dim) + # add a layer to reshape the output to the desired shape + layers.append(nn.Linear(hidden_dims_processed[-1], total_out_dim)) + layers.append(nn.Unflatten(dim=-1, unflattened_size=output_dim)) + + # add last activation function if specified + if last_activation_mod is not None: + layers.append(last_activation_mod) + + # register the layers + for idx, layer in enumerate(layers): + self.add_module(f"{idx}", layer) + + def init_weights(self, scales: float | tuple[float]): + """Initialize the weights of the MLP. + + Args: + scales: Scale factor for the weights. + """ + + def get_scale(idx) -> float: + """Get the scale factor for the weights of the MLP. + + Args: + idx: Index of the layer. + """ + return scales[idx] if isinstance(scales, (list, tuple)) else scales + + # initialize the weights + for idx, module in enumerate(self): + if isinstance(module, nn.Linear): + nn.init.orthogonal_(module.weight, gain=get_scale(idx)) + nn.init.zeros_(module.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass of the MLP. + + Args: + x: Input tensor. + """ + for layer in self: + x = layer(x) + return x + + def reset(self, dones=None, hidden_states=None): + pass + + def detach_hidden_states(self, dones=None): + pass diff --git a/rsl_rl/build/lib/rsl_rl/networks/normalization.py b/rsl_rl/build/lib/rsl_rl/networks/normalization.py new file mode 100644 index 0000000..5fd9692 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/networks/normalization.py @@ -0,0 +1,130 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Copyright (c) 2020 Preferred Networks, Inc. + +from __future__ import annotations + +import torch +from torch import nn + + +class EmpiricalNormalization(nn.Module): + """Normalize mean and variance of values based on empirical values.""" + + def __init__(self, shape, eps=1e-2, until=None): + """Initialize EmpiricalNormalization module. + + Args: + shape (int or tuple of int): Shape of input values except batch axis. + eps (float): Small value for stability. + until (int or None): If this arg is specified, the module learns input values until the sum of batch sizes + exceeds it. + + Note: The normalization parameters are computed over the whole batch, not for each environment separately. + """ + super().__init__() + self.eps = eps + self.until = until + self.register_buffer("_mean", torch.zeros(shape).unsqueeze(0)) + self.register_buffer("_var", torch.ones(shape).unsqueeze(0)) + self.register_buffer("_std", torch.ones(shape).unsqueeze(0)) + self.register_buffer("count", torch.tensor(0, dtype=torch.long)) + + @property + def mean(self): + return self._mean.squeeze(0).clone() + + @property + def std(self): + return self._std.squeeze(0).clone() + + def forward(self, x): + """Normalize mean and variance of values based on empirical values.""" + + return (x - self._mean) / (self._std + self.eps) + + @torch.jit.unused + def update(self, x): + """Learn input values without computing the output values of them""" + + if not self.training: + return + if self.until is not None and self.count >= self.until: + return + + count_x = x.shape[0] + self.count += count_x + rate = count_x / self.count + var_x = torch.var(x, dim=0, unbiased=False, keepdim=True) + mean_x = torch.mean(x, dim=0, keepdim=True) + delta_mean = mean_x - self._mean + self._mean += rate * delta_mean + self._var += rate * (var_x - self._var + delta_mean * (mean_x - self._mean)) + self._std = torch.sqrt(self._var) + + @torch.jit.unused + def inverse(self, y): + """De-normalize values based on empirical values.""" + + return y * (self._std + self.eps) + self._mean + + +class EmpiricalDiscountedVariationNormalization(nn.Module): + """Reward normalization from Pathak's large scale study on PPO. + + Reward normalization. Since the reward function is non-stationary, it is useful to normalize + the scale of the rewards so that the value function can learn quickly. We did this by dividing + the rewards by a running estimate of the standard deviation of the sum of discounted rewards. + """ + + def __init__(self, shape, eps=1e-2, gamma=0.99, until=None): + super().__init__() + + self.emp_norm = EmpiricalNormalization(shape, eps, until) + self.disc_avg = _DiscountedAverage(gamma) + + def forward(self, rew): + if self.training: + # update discounted rewards + avg = self.disc_avg.update(rew) + # update moments from discounted rewards + self.emp_norm.update(avg) + + # normalize rewards with the empirical std + if self.emp_norm._std > 0: + return rew / self.emp_norm._std + else: + return rew + + +""" +Helper class. +""" + + +class _DiscountedAverage: + r"""Discounted average of rewards. + + The discounted average is defined as: + + .. math:: + + \bar{R}_t = \gamma \bar{R}_{t-1} + r_t + + Args: + gamma (float): Discount factor. + """ + + def __init__(self, gamma): + self.avg = None + self.gamma = gamma + + def update(self, rew: torch.Tensor) -> torch.Tensor: + if self.avg is None: + self.avg = rew + else: + self.avg = self.avg * self.gamma + rew + return self.avg diff --git a/rsl_rl/build/lib/rsl_rl/runners/__init__.py b/rsl_rl/build/lib/rsl_rl/runners/__init__.py new file mode 100644 index 0000000..61f1682 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/runners/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Implementation of runners for environment-agent interaction.""" + +from .on_policy_runner import OnPolicyRunner # isort:skip +from .distillation_runner import DistillationRunner +from .amp_on_policy_runner import AMPOnPolicyRunner + +__all__ = ["OnPolicyRunner", "DistillationRunner", "AMPOnPolicyRunner"] diff --git a/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py b/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py new file mode 100644 index 0000000..c0b9b9e --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py @@ -0,0 +1,521 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +import statistics +import time +import torch +import warnings +from collections import deque + +import rsl_rl +from rsl_rl.algorithms import AMP_PPO +from rsl_rl.env import VecEnv +from rsl_rl.modules import ActorCritic, ActorCriticRecurrent,DiscriminatorMulti, resolve_rnd_config, resolve_symmetry_config +from rsl_rl.utils import resolve_obs_groups, store_code_state, Normalizer, G1_AMPLoader + + +class AMPOnPolicyRunner: + """On-policy runner for training and evaluation of actor-critic methods.""" + + def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): + self.cfg = train_cfg + self.alg_cfg = train_cfg["algorithm"] + self.policy_cfg = train_cfg["policy"] + self.device = device + self.env = env + + # check if multi-gpu is enabled + self._configure_multi_gpu() + + # store training configuration + self.num_steps_per_env = self.cfg["num_steps_per_env"] + self.save_interval = self.cfg["save_interval"] + + # query observations from environment for algorithm construction + obs = self.env.get_observations() + default_sets = ["critic"] + if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None: + default_sets.append("rnd_state") + self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets) + + self.amp_data = G1_AMPLoader( + device, + time_between_frames=1/50.0, + preload_transitions=True, + num_preload_transitions=train_cfg["amp_num_preload_transitions"], + motion_files=train_cfg["amp_motion_files"], + num_frames=train_cfg['amp_num_frames'] + ) + + self.amp_observation_dim = self.amp_data.observation_dim if self.cfg["amp_num_obs"] == 0 else self.cfg["amp_num_obs"] + self.amp_num_frames = 0 if self.cfg["amp_num_frames"] == 0 else self.cfg["amp_num_frames"] + self.amp_normalizer = Normalizer(self.amp_observation_dim) + self.discriminator = DiscriminatorMulti( + self.amp_observation_dim, + train_cfg["amp_reward_coef"], + train_cfg["amp_discr_hidden_dims"], + device, + train_cfg["amp_num_frames"], + train_cfg["amp_task_reward_lerp"], + train_cfg['use_lerp'], + ).to(self.device) + + + # create the algorithm + self.alg = self._construct_algorithm(obs) + + # Decide whether to disable logging + # We only log from the process with rank 0 (main process) + self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 + + # Logging + self.log_dir = log_dir + self.writer = None + self.tot_timesteps = 0 + self.tot_time = 0 + self.current_learning_iteration = 0 + self.git_status_repos = [rsl_rl.__file__] + + def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 + # initialize writer + self._prepare_logging_writer() + + # randomize initial episode lengths (for exploration) + if init_at_random_ep_len: + self.env.episode_length_buf = torch.randint_like( + self.env.episode_length_buf, high=int(self.env.max_episode_length) + ) + + # start learning + obs = self.env.get_observations().to(self.device) + self.train_mode() # switch to train mode (for dropout for example) + amp_obs = self.env.get_amp_observations() + amp_obs = amp_obs.to(self.device) + if self.amp_num_frames != 0: + self.amp_obs_frames = torch.zeros(size=(self.env.num_envs, self.amp_num_frames, self.amp_observation_dim), device=self.device) + self.amp_obs_frames = torch.concat((self.amp_obs_frames[:, 1:], amp_obs.unsqueeze(1)), dim=1) + + # Book keeping + ep_infos = [] + rewbuffer = deque(maxlen=100) + lenbuffer = deque(maxlen=100) + step_discrewbuffer = deque(maxlen=100) + + cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_single_step_disc_rew = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + # create buffers for logging extrinsic and intrinsic rewards + if self.alg.rnd: + erewbuffer = deque(maxlen=100) + irewbuffer = deque(maxlen=100) + cur_ereward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_ireward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + # Ensure all parameters are in-synced + if self.is_distributed: + print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") + self.alg.broadcast_parameters() + + # Start training + start_iter = self.current_learning_iteration + tot_iter = start_iter + num_learning_iterations + for it in range(start_iter, tot_iter): + start = time.time() + # Rollout + with torch.inference_mode(): + for _ in range(self.num_steps_per_env): + # Sample actions + actions = self.alg.act(obs,amp_obs) + # Step the environment + obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) + # Move to device + obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) + + next_amp_obs = self.env.get_amp_observations() + next_amp_obs = next_amp_obs.to(self.device) + next_amp_obs_with_term = torch.clone(next_amp_obs) + + reset_env_ids = self.env.reset_env_ids + terminal_amp_states = self.env.get_amp_observations()[reset_env_ids] + next_amp_obs_with_term[reset_env_ids] = terminal_amp_states + self.amp_obs_frames = torch.concat((self.amp_obs_frames[:, 1:], next_amp_obs_with_term.unsqueeze(1)), dim=1) + + amp_reward = torch.zeros(self.env.num_envs, device=obs.device) + + mask = self.env.contact_phase[:, 0] == 1.0 + if mask.any(): + rewards[mask], logit, disc_reward = self.alg.discriminator.predict_amp_reward( + self.amp_obs_frames[mask], rewards[mask], normalizer=self.alg.amp_normalizer + ) + amp_reward[mask] += disc_reward + + # process the step + self.alg.process_env_step(obs, rewards, dones, extras, next_amp_obs_with_term, self.amp_obs_frames) + self.amp_obs_frames[reset_env_ids] = 0 + + amp_obs = torch.clone(next_amp_obs) + # Extract intrinsic rewards (only for logging) + intrinsic_rewards = self.alg.intrinsic_rewards if self.alg.rnd else None + # book keeping + if self.log_dir is not None: + if "episode" in extras: + ep_infos.append(extras["episode"]) + elif "log" in extras: + ep_infos.append(extras["log"]) + # Update rewards + if self.alg.rnd: + cur_ereward_sum += rewards + cur_ireward_sum += intrinsic_rewards # type: ignore + cur_reward_sum += rewards + intrinsic_rewards + else: + cur_reward_sum += rewards + # Update episode length + cur_episode_length += 1 + cur_single_step_disc_rew += amp_reward + # Clear data for completed episodes + # -- common + new_ids = (dones > 0).nonzero(as_tuple=False) + rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) + lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) + cur_reward_sum[new_ids] = 0 + cur_episode_length[new_ids] = 0 + to_extend_disc = (cur_single_step_disc_rew[new_ids] / self.env.max_episode_length_s)[:, 0].cpu().numpy() + step_discrewbuffer.extend(to_extend_disc.tolist()) + cur_single_step_disc_rew[new_ids] = 0 + # -- intrinsic and extrinsic rewards + if self.alg.rnd: + erewbuffer.extend(cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist()) + irewbuffer.extend(cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist()) + cur_ereward_sum[new_ids] = 0 + cur_ireward_sum[new_ids] = 0 + + stop = time.time() + collection_time = stop - start + start = stop + + # compute returns + self.alg.compute_returns(obs) + + # update policy + loss_dict = self.alg.update() + + stop = time.time() + learn_time = stop - start + self.current_learning_iteration = it + # log info + if self.log_dir is not None and not self.disable_logs: + # Log information + self.log(locals()) + # Save model + if it % self.save_interval == 0: + self.save(os.path.join(self.log_dir, f"model_{it}.pt")) + + # Clear episode infos + ep_infos.clear() + # Save code state + if it == start_iter and not self.disable_logs: + # obtain all the diff files + git_file_paths = store_code_state(self.log_dir, self.git_status_repos) + # if possible store them to wandb + if self.logger_type in ["wandb", "neptune"] and git_file_paths: + for path in git_file_paths: + self.writer.save_file(path) + + # Save the final model after training + if self.log_dir is not None and not self.disable_logs: + self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) + + def log(self, locs: dict, width: int = 80, pad: int = 35): + # Compute the collection size + collection_size = self.num_steps_per_env * self.env.num_envs * self.gpu_world_size + # Update total time-steps and time + self.tot_timesteps += collection_size + self.tot_time += locs["collection_time"] + locs["learn_time"] + iteration_time = locs["collection_time"] + locs["learn_time"] + + # -- Episode info + ep_string = "" + if locs["ep_infos"]: + for key in locs["ep_infos"][0]: + infotensor = torch.tensor([], device=self.device) + for ep_info in locs["ep_infos"]: + # handle scalar and zero dimensional tensor infos + if key not in ep_info: + continue + if not isinstance(ep_info[key], torch.Tensor): + ep_info[key] = torch.Tensor([ep_info[key]]) + if len(ep_info[key].shape) == 0: + ep_info[key] = ep_info[key].unsqueeze(0) + infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) + value = torch.mean(infotensor) + # log to logger and terminal + if "/" in key: + self.writer.add_scalar(key, value, locs["it"]) + ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + else: + self.writer.add_scalar("Episode/" + key, value, locs["it"]) + ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" + + mean_std = self.alg.policy.action_std.mean() + fps = int(collection_size / (locs["collection_time"] + locs["learn_time"])) + + # -- Losses + for key, value in locs["loss_dict"].items(): + self.writer.add_scalar(f"Loss/{key}", value, locs["it"]) + self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) + + # -- Policy + self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) + + # -- Performance + self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) + self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) + self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) + + # -- Training + if len(locs["rewbuffer"]) > 0: + # separate logging for intrinsic and extrinsic rewards + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(locs["erewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(locs["irewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/weight", self.alg.rnd.weight, locs["it"]) + # everything else + self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) + self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) + self.writer.add_scalar('Train/mean_step_disc_reward', statistics.mean(locs['step_discrewbuffer']), locs['it']) + if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging + self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) + self.writer.add_scalar( + "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time + ) + + str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " + + if len(locs["rewbuffer"]) > 0: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + f"""{'Step disc reward:':>{pad}} {statistics.mean(locs['step_discrewbuffer']):.2f}\n""" + ) + # -- Losses + for key, value in locs["loss_dict"].items(): + log_string += f"""{f'Mean {key} loss:':>{pad}} {value:.4f}\n""" + # -- Rewards + if hasattr(self.alg, "rnd") and self.alg.rnd: + log_string += ( + f"""{'Mean extrinsic reward:':>{pad}} {statistics.mean(locs['erewbuffer']):.2f}\n""" + f"""{'Mean intrinsic reward:':>{pad}} {statistics.mean(locs['irewbuffer']):.2f}\n""" + ) + log_string += f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" + # -- episode info + log_string += f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" + else: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + ) + for key, value in locs["loss_dict"].items(): + log_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + + log_string += ep_string + log_string += ( + f"""{'-' * width}\n""" + f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" + f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" + f"""{'Time elapsed:':>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n""" + f"""{'ETA:':>{pad}} {time.strftime( + "%H:%M:%S", + time.gmtime( + self.tot_time / (locs['it'] - locs['start_iter'] + 1) + * (locs['start_iter'] + locs['num_learning_iterations'] - locs['it']) + ) + )}\n""" + ) + print(log_string) + + def save(self, path: str, infos=None): + # -- Save model + saved_dict = { + "model_state_dict": self.alg.policy.state_dict(), + "optimizer_state_dict": self.alg.optimizer.state_dict(), + "iter": self.current_learning_iteration, + "infos": infos, + } + # -- Save RND model if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict() + saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict() + torch.save(saved_dict, path) + + # upload model to external logging service + if self.logger_type in ["neptune", "wandb"] and not self.disable_logs: + self.writer.save_model(path, self.current_learning_iteration) + + def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None): + loaded_dict = torch.load(path, weights_only=False, map_location=map_location) + # -- Load model + resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"]) + # -- Load RND model if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"]) + # -- load optimizer if used + if load_optimizer and resumed_training: + # -- algorithm optimizer + self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) + # -- RND optimizer if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"]) + # -- load current learning iteration + if resumed_training: + self.current_learning_iteration = loaded_dict["iter"] + return loaded_dict["infos"] + + def get_inference_policy(self, device=None): + self.eval_mode() # switch to evaluation mode (dropout for example) + if device is not None: + self.alg.policy.to(device) + return self.alg.policy.act_inference + + def train_mode(self): + # -- PPO + self.alg.policy.train() + self.alg.discriminator.train() + # -- RND + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.train() + + def eval_mode(self): + # -- PPO + self.alg.policy.eval() + self.alg.discriminator.eval() + # -- RND + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.eval() + + def add_git_repo_to_log(self, repo_file_path): + self.git_status_repos.append(repo_file_path) + + """ + Helper functions. + """ + + def _configure_multi_gpu(self): + """Configure multi-gpu training.""" + # check if distributed training is enabled + self.gpu_world_size = int(os.getenv("WORLD_SIZE", "1")) + self.is_distributed = self.gpu_world_size > 1 + + # if not distributed training, set local and global rank to 0 and return + if not self.is_distributed: + self.gpu_local_rank = 0 + self.gpu_global_rank = 0 + self.multi_gpu_cfg = None + return + + # get rank and world size + self.gpu_local_rank = int(os.getenv("LOCAL_RANK", "0")) + self.gpu_global_rank = int(os.getenv("RANK", "0")) + + # make a configuration dictionary + self.multi_gpu_cfg = { + "global_rank": self.gpu_global_rank, # rank of the main process + "local_rank": self.gpu_local_rank, # rank of the current process + "world_size": self.gpu_world_size, # total number of processes + } + + # check if user has device specified for local rank + if self.device != f"cuda:{self.gpu_local_rank}": + raise ValueError( + f"Device '{self.device}' does not match expected device for local rank '{self.gpu_local_rank}'." + ) + # validate multi-gpu configuration + if self.gpu_local_rank >= self.gpu_world_size: + raise ValueError( + f"Local rank '{self.gpu_local_rank}' is greater than or equal to world size '{self.gpu_world_size}'." + ) + if self.gpu_global_rank >= self.gpu_world_size: + raise ValueError( + f"Global rank '{self.gpu_global_rank}' is greater than or equal to world size '{self.gpu_world_size}'." + ) + + # initialize torch distributed + torch.distributed.init_process_group(backend="nccl", rank=self.gpu_global_rank, world_size=self.gpu_world_size) + # set device to the local rank + torch.cuda.set_device(self.gpu_local_rank) + + def _construct_algorithm(self, obs) -> AMP_PPO: + """Construct the actor-critic algorithm.""" + # resolve RND config + self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env) + + # resolve symmetry config + self.alg_cfg = resolve_symmetry_config(self.alg_cfg, self.env) + + # resolve deprecated normalization config + if self.cfg.get("empirical_normalization") is not None: + warnings.warn( + "The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and " + "`critic_obs_normalization` as part of the `policy` configuration instead.", + DeprecationWarning, + ) + if self.policy_cfg.get("actor_obs_normalization") is None: + self.policy_cfg["actor_obs_normalization"] = self.cfg["empirical_normalization"] + if self.policy_cfg.get("critic_obs_normalization") is None: + self.policy_cfg["critic_obs_normalization"] = self.cfg["empirical_normalization"] + + # initialize the actor-critic + actor_critic_class = eval(self.policy_cfg.pop("class_name")) + actor_critic: ActorCritic | ActorCriticRecurrent = actor_critic_class( + obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg + ).to(self.device) + + # initialize the algorithm + alg_class = eval(self.alg_cfg.pop("class_name")) + + alg: AMP_PPO = alg_class(actor_critic, self.discriminator, self.amp_data, self.amp_normalizer, self.amp_num_frames, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg) + + # initialize the storage + alg.init_storage( + "rl", + self.env.num_envs, + self.num_steps_per_env, + obs, + [self.env.num_actions], + ) + + return alg + + def _prepare_logging_writer(self): + """Prepares the logging writers.""" + if self.log_dir is not None and self.writer is None and not self.disable_logs: + # Launch either Tensorboard or Neptune & Tensorboard summary writer(s), default: Tensorboard. + self.logger_type = self.cfg.get("logger", "tensorboard") + self.logger_type = self.logger_type.lower() + + if self.logger_type == "neptune": + from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter + + self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "wandb": + from rsl_rl.utils.wandb_utils import WandbSummaryWriter + + self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "tensorboard": + from torch.utils.tensorboard import SummaryWriter + + self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) + else: + raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.") diff --git a/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py b/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py new file mode 100644 index 0000000..9cc6a8b --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py @@ -0,0 +1,179 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +import time +import torch +from collections import deque + +import rsl_rl +from rsl_rl.algorithms import Distillation +from rsl_rl.env import VecEnv +from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent +from rsl_rl.runners import OnPolicyRunner +from rsl_rl.utils import resolve_obs_groups, store_code_state + + +class DistillationRunner(OnPolicyRunner): + """On-policy runner for training and evaluation of teacher-student training.""" + + def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): + self.cfg = train_cfg + self.alg_cfg = train_cfg["algorithm"] + self.policy_cfg = train_cfg["policy"] + self.device = device + self.env = env + + # check if multi-gpu is enabled + self._configure_multi_gpu() + + # store training configuration + self.num_steps_per_env = self.cfg["num_steps_per_env"] + self.save_interval = self.cfg["save_interval"] + + # query observations from environment for algorithm construction + obs = self.env.get_observations() + self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets=["teacher"]) + + # create the algorithm + self.alg = self._construct_algorithm(obs) + + # Decide whether to disable logging + # We only log from the process with rank 0 (main process) + self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 + + # Logging + self.log_dir = log_dir + self.writer = None + self.tot_timesteps = 0 + self.tot_time = 0 + self.current_learning_iteration = 0 + self.git_status_repos = [rsl_rl.__file__] + + def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 + # initialize writer + self._prepare_logging_writer() + # check if teacher is loaded + if not self.alg.policy.loaded_teacher: + raise ValueError("Teacher model parameters not loaded. Please load a teacher model to distill.") + + # randomize initial episode lengths (for exploration) + if init_at_random_ep_len: + self.env.episode_length_buf = torch.randint_like( + self.env.episode_length_buf, high=int(self.env.max_episode_length) + ) + + # start learning + obs = self.env.get_observations().to(self.device) + self.train_mode() # switch to train mode (for dropout for example) + + # Book keeping + ep_infos = [] + rewbuffer = deque(maxlen=100) + lenbuffer = deque(maxlen=100) + cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + # Ensure all parameters are in-synced + if self.is_distributed: + print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") + self.alg.broadcast_parameters() + + # Start training + start_iter = self.current_learning_iteration + tot_iter = start_iter + num_learning_iterations + for it in range(start_iter, tot_iter): + start = time.time() + # Rollout + with torch.inference_mode(): + for _ in range(self.num_steps_per_env): + # Sample actions + actions = self.alg.act(obs) + # Step the environment + obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) + # Move to device + obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) + # process the step + self.alg.process_env_step(obs, rewards, dones, extras) + # book keeping + if self.log_dir is not None: + if "episode" in extras: + ep_infos.append(extras["episode"]) + elif "log" in extras: + ep_infos.append(extras["log"]) + # Update rewards + cur_reward_sum += rewards + # Update episode length + cur_episode_length += 1 + # Clear data for completed episodes + new_ids = (dones > 0).nonzero(as_tuple=False) + rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) + lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) + cur_reward_sum[new_ids] = 0 + cur_episode_length[new_ids] = 0 + + stop = time.time() + collection_time = stop - start + start = stop + + # update policy + loss_dict = self.alg.update() + + stop = time.time() + learn_time = stop - start + self.current_learning_iteration = it + # log info + if self.log_dir is not None and not self.disable_logs: + # Log information + self.log(locals()) + # Save model + if it % self.save_interval == 0: + self.save(os.path.join(self.log_dir, f"model_{it}.pt")) + + # Clear episode infos + ep_infos.clear() + # Save code state + if it == start_iter and not self.disable_logs: + # obtain all the diff files + git_file_paths = store_code_state(self.log_dir, self.git_status_repos) + # if possible store them to wandb + if self.logger_type in ["wandb", "neptune"] and git_file_paths: + for path in git_file_paths: + self.writer.save_file(path) + + # Save the final model after training + if self.log_dir is not None and not self.disable_logs: + self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) + + """ + Helper methods. + """ + + def _construct_algorithm(self, obs) -> Distillation: + """Construct the distillation algorithm.""" + # initialize the actor-critic + student_teacher_class = eval(self.policy_cfg.pop("class_name")) + student_teacher: StudentTeacher | StudentTeacherRecurrent = student_teacher_class( + obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg + ).to(self.device) + + # initialize the algorithm + alg_class = eval(self.alg_cfg.pop("class_name")) + alg: Distillation = alg_class( + student_teacher, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg + ) + + # initialize the storage + alg.init_storage( + "distillation", + self.env.num_envs, + self.num_steps_per_env, + obs, + [self.env.num_actions], + ) + + return alg diff --git a/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py b/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py new file mode 100644 index 0000000..36f11f3 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py @@ -0,0 +1,460 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +import statistics +import time +import torch +import warnings +from collections import deque + +import rsl_rl +from rsl_rl.algorithms import PPO +from rsl_rl.env import VecEnv +from rsl_rl.modules import ActorCritic, ActorCriticRecurrent, resolve_rnd_config, resolve_symmetry_config +from rsl_rl.utils import resolve_obs_groups, store_code_state + + +class OnPolicyRunner: + """On-policy runner for training and evaluation of actor-critic methods.""" + + def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): + self.cfg = train_cfg + self.alg_cfg = train_cfg["algorithm"] + self.policy_cfg = train_cfg["policy"] + self.device = device + self.env = env + + # check if multi-gpu is enabled + self._configure_multi_gpu() + + # store training configuration + self.num_steps_per_env = self.cfg["num_steps_per_env"] + self.save_interval = self.cfg["save_interval"] + + # query observations from environment for algorithm construction + obs = self.env.get_observations() + default_sets = ["critic"] + if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None: + default_sets.append("rnd_state") + self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets) + + # create the algorithm + self.alg = self._construct_algorithm(obs) + + # Decide whether to disable logging + # We only log from the process with rank 0 (main process) + self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 + + # Logging + self.log_dir = log_dir + self.writer = None + self.tot_timesteps = 0 + self.tot_time = 0 + self.current_learning_iteration = 0 + self.git_status_repos = [rsl_rl.__file__] + + def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 + # initialize writer + self._prepare_logging_writer() + + # randomize initial episode lengths (for exploration) + if init_at_random_ep_len: + self.env.episode_length_buf = torch.randint_like( + self.env.episode_length_buf, high=int(self.env.max_episode_length) + ) + + # start learning + obs = self.env.get_observations().to(self.device) + self.train_mode() # switch to train mode (for dropout for example) + + # Book keeping + ep_infos = [] + rewbuffer = deque(maxlen=100) + lenbuffer = deque(maxlen=100) + cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + # create buffers for logging extrinsic and intrinsic rewards + if self.alg.rnd: + erewbuffer = deque(maxlen=100) + irewbuffer = deque(maxlen=100) + cur_ereward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_ireward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + # Ensure all parameters are in-synced + if self.is_distributed: + print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") + self.alg.broadcast_parameters() + + # Start training + start_iter = self.current_learning_iteration + tot_iter = start_iter + num_learning_iterations + for it in range(start_iter, tot_iter): + start = time.time() + # Rollout + with torch.inference_mode(): + for _ in range(self.num_steps_per_env): + # Sample actions + actions = self.alg.act(obs) + # Step the environment + obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) + # Move to device + obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) + # process the step + self.alg.process_env_step(obs, rewards, dones, extras) + # Extract intrinsic rewards (only for logging) + intrinsic_rewards = self.alg.intrinsic_rewards if self.alg.rnd else None + # book keeping + if self.log_dir is not None: + if "episode" in extras: + ep_infos.append(extras["episode"]) + elif "log" in extras: + ep_infos.append(extras["log"]) + # Update rewards + if self.alg.rnd: + cur_ereward_sum += rewards + cur_ireward_sum += intrinsic_rewards # type: ignore + cur_reward_sum += rewards + intrinsic_rewards + else: + cur_reward_sum += rewards + # Update episode length + cur_episode_length += 1 + # Clear data for completed episodes + # -- common + new_ids = (dones > 0).nonzero(as_tuple=False) + rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) + lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) + cur_reward_sum[new_ids] = 0 + cur_episode_length[new_ids] = 0 + # -- intrinsic and extrinsic rewards + if self.alg.rnd: + erewbuffer.extend(cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist()) + irewbuffer.extend(cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist()) + cur_ereward_sum[new_ids] = 0 + cur_ireward_sum[new_ids] = 0 + + stop = time.time() + collection_time = stop - start + start = stop + + # compute returns + self.alg.compute_returns(obs) + + # update policy + loss_dict = self.alg.update() + + stop = time.time() + learn_time = stop - start + self.current_learning_iteration = it + # log info + if self.log_dir is not None and not self.disable_logs: + # Log information + self.log(locals()) + # Save model + if it % self.save_interval == 0: + self.save(os.path.join(self.log_dir, f"model_{it}.pt")) + + # Clear episode infos + ep_infos.clear() + # Save code state + if it == start_iter and not self.disable_logs: + # obtain all the diff files + git_file_paths = store_code_state(self.log_dir, self.git_status_repos) + # if possible store them to wandb + if self.logger_type in ["wandb", "neptune"] and git_file_paths: + for path in git_file_paths: + self.writer.save_file(path) + + # Save the final model after training + if self.log_dir is not None and not self.disable_logs: + self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) + + def log(self, locs: dict, width: int = 80, pad: int = 35): + # Compute the collection size + collection_size = self.num_steps_per_env * self.env.num_envs * self.gpu_world_size + # Update total time-steps and time + self.tot_timesteps += collection_size + self.tot_time += locs["collection_time"] + locs["learn_time"] + iteration_time = locs["collection_time"] + locs["learn_time"] + + # -- Episode info + ep_string = "" + if locs["ep_infos"]: + for key in locs["ep_infos"][0]: + infotensor = torch.tensor([], device=self.device) + for ep_info in locs["ep_infos"]: + # handle scalar and zero dimensional tensor infos + if key not in ep_info: + continue + if not isinstance(ep_info[key], torch.Tensor): + ep_info[key] = torch.Tensor([ep_info[key]]) + if len(ep_info[key].shape) == 0: + ep_info[key] = ep_info[key].unsqueeze(0) + infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) + value = torch.mean(infotensor) + # log to logger and terminal + if "/" in key: + self.writer.add_scalar(key, value, locs["it"]) + ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + else: + self.writer.add_scalar("Episode/" + key, value, locs["it"]) + ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" + + mean_std = self.alg.policy.action_std.mean() + fps = int(collection_size / (locs["collection_time"] + locs["learn_time"])) + + # -- Losses + for key, value in locs["loss_dict"].items(): + self.writer.add_scalar(f"Loss/{key}", value, locs["it"]) + self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) + + # -- Policy + self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) + + # -- Performance + self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) + self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) + self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) + + # -- Training + if len(locs["rewbuffer"]) > 0: + # separate logging for intrinsic and extrinsic rewards + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(locs["erewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(locs["irewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/weight", self.alg.rnd.weight, locs["it"]) + # everything else + self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) + self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) + if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging + self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) + self.writer.add_scalar( + "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time + ) + + str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " + + if len(locs["rewbuffer"]) > 0: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + ) + # -- Losses + for key, value in locs["loss_dict"].items(): + log_string += f"""{f'Mean {key} loss:':>{pad}} {value:.4f}\n""" + # -- Rewards + if hasattr(self.alg, "rnd") and self.alg.rnd: + log_string += ( + f"""{'Mean extrinsic reward:':>{pad}} {statistics.mean(locs['erewbuffer']):.2f}\n""" + f"""{'Mean intrinsic reward:':>{pad}} {statistics.mean(locs['irewbuffer']):.2f}\n""" + ) + log_string += f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" + # -- episode info + log_string += f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" + else: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + ) + for key, value in locs["loss_dict"].items(): + log_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + + log_string += ep_string + log_string += ( + f"""{'-' * width}\n""" + f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" + f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" + f"""{'Time elapsed:':>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n""" + f"""{'ETA:':>{pad}} {time.strftime( + "%H:%M:%S", + time.gmtime( + self.tot_time / (locs['it'] - locs['start_iter'] + 1) + * (locs['start_iter'] + locs['num_learning_iterations'] - locs['it']) + ) + )}\n""" + ) + print(log_string) + + def save(self, path: str, infos=None): + # -- Save model + saved_dict = { + "model_state_dict": self.alg.policy.state_dict(), + "optimizer_state_dict": self.alg.optimizer.state_dict(), + "iter": self.current_learning_iteration, + "infos": infos, + } + # -- Save RND model if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict() + saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict() + torch.save(saved_dict, path) + + # upload model to external logging service + if self.logger_type in ["neptune", "wandb"] and not self.disable_logs: + self.writer.save_model(path, self.current_learning_iteration) + + def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None): + loaded_dict = torch.load(path, weights_only=False, map_location=map_location) + # -- Load model + resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"]) + # -- Load RND model if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"]) + # -- load optimizer if used + if load_optimizer and resumed_training: + # -- algorithm optimizer + self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) + # -- RND optimizer if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"]) + # -- load current learning iteration + if resumed_training: + self.current_learning_iteration = loaded_dict["iter"] + return loaded_dict["infos"] + + def get_inference_policy(self, device=None): + self.eval_mode() # switch to evaluation mode (dropout for example) + if device is not None: + self.alg.policy.to(device) + return self.alg.policy.act_inference + + def train_mode(self): + # -- PPO + self.alg.policy.train() + # -- RND + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.train() + + def eval_mode(self): + # -- PPO + self.alg.policy.eval() + # -- RND + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.eval() + + def add_git_repo_to_log(self, repo_file_path): + self.git_status_repos.append(repo_file_path) + + """ + Helper functions. + """ + + def _configure_multi_gpu(self): + """Configure multi-gpu training.""" + # check if distributed training is enabled + self.gpu_world_size = int(os.getenv("WORLD_SIZE", "1")) + self.is_distributed = self.gpu_world_size > 1 + + # if not distributed training, set local and global rank to 0 and return + if not self.is_distributed: + self.gpu_local_rank = 0 + self.gpu_global_rank = 0 + self.multi_gpu_cfg = None + return + + # get rank and world size + self.gpu_local_rank = int(os.getenv("LOCAL_RANK", "0")) + self.gpu_global_rank = int(os.getenv("RANK", "0")) + + # make a configuration dictionary + self.multi_gpu_cfg = { + "global_rank": self.gpu_global_rank, # rank of the main process + "local_rank": self.gpu_local_rank, # rank of the current process + "world_size": self.gpu_world_size, # total number of processes + } + + # check if user has device specified for local rank + if self.device != f"cuda:{self.gpu_local_rank}": + raise ValueError( + f"Device '{self.device}' does not match expected device for local rank '{self.gpu_local_rank}'." + ) + # validate multi-gpu configuration + if self.gpu_local_rank >= self.gpu_world_size: + raise ValueError( + f"Local rank '{self.gpu_local_rank}' is greater than or equal to world size '{self.gpu_world_size}'." + ) + if self.gpu_global_rank >= self.gpu_world_size: + raise ValueError( + f"Global rank '{self.gpu_global_rank}' is greater than or equal to world size '{self.gpu_world_size}'." + ) + + # initialize torch distributed + torch.distributed.init_process_group(backend="nccl", rank=self.gpu_global_rank, world_size=self.gpu_world_size) + # set device to the local rank + torch.cuda.set_device(self.gpu_local_rank) + + def _construct_algorithm(self, obs) -> PPO: + """Construct the actor-critic algorithm.""" + # resolve RND config + self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env) + + # resolve symmetry config + self.alg_cfg = resolve_symmetry_config(self.alg_cfg, self.env) + + # resolve deprecated normalization config + if self.cfg.get("empirical_normalization") is not None: + warnings.warn( + "The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and " + "`critic_obs_normalization` as part of the `policy` configuration instead.", + DeprecationWarning, + ) + if self.policy_cfg.get("actor_obs_normalization") is None: + self.policy_cfg["actor_obs_normalization"] = self.cfg["empirical_normalization"] + if self.policy_cfg.get("critic_obs_normalization") is None: + self.policy_cfg["critic_obs_normalization"] = self.cfg["empirical_normalization"] + + # initialize the actor-critic + actor_critic_class = eval(self.policy_cfg.pop("class_name")) + actor_critic: ActorCritic | ActorCriticRecurrent = actor_critic_class( + obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg + ).to(self.device) + + # initialize the algorithm + alg_class = eval(self.alg_cfg.pop("class_name")) + alg: PPO = alg_class(actor_critic, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg) + + # initialize the storage + alg.init_storage( + "rl", + self.env.num_envs, + self.num_steps_per_env, + obs, + [self.env.num_actions], + ) + + return alg + + def _prepare_logging_writer(self): + """Prepares the logging writers.""" + if self.log_dir is not None and self.writer is None and not self.disable_logs: + # Launch either Tensorboard or Neptune & Tensorboard summary writer(s), default: Tensorboard. + self.logger_type = self.cfg.get("logger", "tensorboard") + self.logger_type = self.logger_type.lower() + + if self.logger_type == "neptune": + from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter + + self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "wandb": + from rsl_rl.utils.wandb_utils import WandbSummaryWriter + + self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "tensorboard": + from torch.utils.tensorboard import SummaryWriter + + self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) + else: + raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.") diff --git a/rsl_rl/build/lib/rsl_rl/storage/__init__.py b/rsl_rl/build/lib/rsl_rl/storage/__init__.py new file mode 100644 index 0000000..1624330 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/storage/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Implementation of transitions storage for RL-agent.""" + +from .rollout_storage import RolloutStorage +from .replay_buffer_multi import ReplayBufferMulti +__all__ = ["RolloutStorage", "ReplayBufferMulti"] diff --git a/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py b/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py new file mode 100644 index 0000000..6462b8e --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py @@ -0,0 +1,38 @@ +import torch +import numpy as np + + +class ReplayBufferMulti: + """Fixed-size buffer to store experience tuples.""" + + def __init__(self, obs_dim, buffer_size, num_amp_frames, device): + """Initialize a ReplayBuffer object. + Arguments: + buffer_size (int): maximum size of buffer + """ + self.states = torch.zeros(buffer_size, num_amp_frames, obs_dim).to(device) + self.num_amp_frames = num_amp_frames + self.buffer_size = buffer_size + self.device = device + + self.step = 0 + self.num_samples = 0 + + def insert(self, states): + """Add new states to memory.""" + num_states = states.shape[0] + start_idx = self.step + end_idx = self.step + num_states + if end_idx > self.buffer_size: + self.states[self.step:self.buffer_size] = states[:self.buffer_size - self.step] + self.states[:end_idx - self.buffer_size] = states[self.buffer_size - self.step:] + else: + self.states[start_idx:end_idx] = states + + self.num_samples = min(self.buffer_size, max(end_idx, self.num_samples)) + self.step = (self.step + num_states) % self.buffer_size + + def feed_forward_generator(self, num_mini_batch, mini_batch_size): + for _ in range(num_mini_batch): + sample_idxs = np.random.choice(self.num_samples, size=mini_batch_size) + yield (self.states[sample_idxs].to(self.device)) diff --git a/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py b/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py new file mode 100644 index 0000000..e9309b3 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py @@ -0,0 +1,260 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from tensordict import TensorDict + +from rsl_rl.utils import split_and_pad_trajectories + + +class RolloutStorage: + class Transition: + def __init__(self): + self.observations = None + self.actions = None + self.privileged_actions = None + self.rewards = None + self.dones = None + self.values = None + self.actions_log_prob = None + self.action_mean = None + self.action_sigma = None + self.hidden_states = None + + def clear(self): + self.__init__() + + def __init__( + self, + training_type, + num_envs, + num_transitions_per_env, + obs, + actions_shape, + device="cpu", + ): + # store inputs + self.training_type = training_type + self.device = device + self.num_transitions_per_env = num_transitions_per_env + self.num_envs = num_envs + self.actions_shape = actions_shape + + # Core + self.observations = TensorDict( + {key: torch.zeros(num_transitions_per_env, *value.shape, device=device) for key, value in obs.items()}, + batch_size=[num_transitions_per_env, num_envs], + device=self.device, + ) + self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte() + + # for distillation + if training_type == "distillation": + self.privileged_actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + + # for reinforcement learning + if training_type == "rl": + self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + + # For RNN networks + self.saved_hidden_states_a = None + self.saved_hidden_states_c = None + + # counter for the number of transitions stored + self.step = 0 + + def add_transitions(self, transition: Transition): + # check if the transition is valid + if self.step >= self.num_transitions_per_env: + raise OverflowError("Rollout buffer overflow! You should call clear() before adding new transitions.") + + # Core + self.observations[self.step].copy_(transition.observations) + self.actions[self.step].copy_(transition.actions) + self.rewards[self.step].copy_(transition.rewards.view(-1, 1)) + self.dones[self.step].copy_(transition.dones.view(-1, 1)) + + # for distillation + if self.training_type == "distillation": + self.privileged_actions[self.step].copy_(transition.privileged_actions) + + # for reinforcement learning + if self.training_type == "rl": + self.values[self.step].copy_(transition.values) + self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1)) + self.mu[self.step].copy_(transition.action_mean) + self.sigma[self.step].copy_(transition.action_sigma) + + # For RNN networks + self._save_hidden_states(transition.hidden_states) + + # increment the counter + self.step += 1 + + def _save_hidden_states(self, hidden_states): + if hidden_states is None or hidden_states == (None, None): + return + # make a tuple out of GRU hidden state sto match the LSTM format + hid_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],) + hid_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],) + # initialize if needed + if self.saved_hidden_states_a is None: + self.saved_hidden_states_a = [ + torch.zeros(self.observations.shape[0], *hid_a[i].shape, device=self.device) for i in range(len(hid_a)) + ] + self.saved_hidden_states_c = [ + torch.zeros(self.observations.shape[0], *hid_c[i].shape, device=self.device) for i in range(len(hid_c)) + ] + # copy the states + for i in range(len(hid_a)): + self.saved_hidden_states_a[i][self.step].copy_(hid_a[i]) + self.saved_hidden_states_c[i][self.step].copy_(hid_c[i]) + + def clear(self): + self.step = 0 + + def compute_returns(self, last_values, gamma, lam, normalize_advantage: bool = True): + advantage = 0 + for step in reversed(range(self.num_transitions_per_env)): + # if we are at the last step, bootstrap the return value + if step == self.num_transitions_per_env - 1: + next_values = last_values + else: + next_values = self.values[step + 1] + # 1 if we are not in a terminal state, 0 otherwise + next_is_not_terminal = 1.0 - self.dones[step].float() + # TD error: r_t + gamma * V(s_{t+1}) - V(s_t) + delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step] + # Advantage: A(s_t, a_t) = delta_t + gamma * lambda * A(s_{t+1}, a_{t+1}) + advantage = delta + next_is_not_terminal * gamma * lam * advantage + # Return: R_t = A(s_t, a_t) + V(s_t) + self.returns[step] = advantage + self.values[step] + + # Compute the advantages + self.advantages = self.returns - self.values + # Normalize the advantages if flag is set + # This is to prevent double normalization (i.e. if per minibatch normalization is used) + if normalize_advantage: + self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8) + + # for distillation + def generator(self): + if self.training_type != "distillation": + raise ValueError("This function is only available for distillation training.") + + for i in range(self.num_transitions_per_env): + yield self.observations[i], self.actions[i], self.privileged_actions[i], self.dones[i] + + # for reinforcement learning with feedforward networks + def mini_batch_generator(self, num_mini_batches, num_epochs=8): + if self.training_type != "rl": + raise ValueError("This function is only available for reinforcement learning training.") + batch_size = self.num_envs * self.num_transitions_per_env + mini_batch_size = batch_size // num_mini_batches + indices = torch.randperm(num_mini_batches * mini_batch_size, requires_grad=False, device=self.device) + + # Core + observations = self.observations.flatten(0, 1) + actions = self.actions.flatten(0, 1) + values = self.values.flatten(0, 1) + returns = self.returns.flatten(0, 1) + + # For PPO + old_actions_log_prob = self.actions_log_prob.flatten(0, 1) + advantages = self.advantages.flatten(0, 1) + old_mu = self.mu.flatten(0, 1) + old_sigma = self.sigma.flatten(0, 1) + + for epoch in range(num_epochs): + for i in range(num_mini_batches): + # Select the indices for the mini-batch + start = i * mini_batch_size + end = (i + 1) * mini_batch_size + batch_idx = indices[start:end] + + # Create the mini-batch + # -- Core + obs_batch = observations[batch_idx] + actions_batch = actions[batch_idx] + + # -- For PPO + target_values_batch = values[batch_idx] + returns_batch = returns[batch_idx] + old_actions_log_prob_batch = old_actions_log_prob[batch_idx] + advantages_batch = advantages[batch_idx] + old_mu_batch = old_mu[batch_idx] + old_sigma_batch = old_sigma[batch_idx] + + # yield the mini-batch + yield obs_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( + None, + None, + ), None + + # for reinfrocement learning with recurrent networks + def recurrent_mini_batch_generator(self, num_mini_batches, num_epochs=8): + if self.training_type != "rl": + raise ValueError("This function is only available for reinforcement learning training.") + padded_obs_trajectories, trajectory_masks = split_and_pad_trajectories(self.observations, self.dones) + + mini_batch_size = self.num_envs // num_mini_batches + for ep in range(num_epochs): + first_traj = 0 + for i in range(num_mini_batches): + start = i * mini_batch_size + stop = (i + 1) * mini_batch_size + + dones = self.dones.squeeze(-1) + last_was_done = torch.zeros_like(dones, dtype=torch.bool) + last_was_done[1:] = dones[:-1] + last_was_done[0] = True + trajectories_batch_size = torch.sum(last_was_done[:, start:stop]) + last_traj = first_traj + trajectories_batch_size + + masks_batch = trajectory_masks[:, first_traj:last_traj] + obs_batch = padded_obs_trajectories[:, first_traj:last_traj] + actions_batch = self.actions[:, start:stop] + old_mu_batch = self.mu[:, start:stop] + old_sigma_batch = self.sigma[:, start:stop] + returns_batch = self.returns[:, start:stop] + advantages_batch = self.advantages[:, start:stop] + values_batch = self.values[:, start:stop] + old_actions_log_prob_batch = self.actions_log_prob[:, start:stop] + + # reshape to [num_envs, time, num layers, hidden dim] (original shape: [time, num_layers, num_envs, hidden_dim]) + # then take only time steps after dones (flattens num envs and time dimensions), + # take a batch of trajectories and finally reshape back to [num_layers, batch, hidden_dim] + last_was_done = last_was_done.permute(1, 0) + hid_a_batch = [ + saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] + .transpose(1, 0) + .contiguous() + for saved_hidden_states in self.saved_hidden_states_a + ] + hid_c_batch = [ + saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] + .transpose(1, 0) + .contiguous() + for saved_hidden_states in self.saved_hidden_states_c + ] + # remove the tuple for GRU + hid_a_batch = hid_a_batch[0] if len(hid_a_batch) == 1 else hid_a_batch + hid_c_batch = hid_c_batch[0] if len(hid_c_batch) == 1 else hid_c_batch + + yield obs_batch, actions_batch, values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( + hid_a_batch, + hid_c_batch, + ), masks_batch + + first_traj = last_traj diff --git a/rsl_rl/build/lib/rsl_rl/utils/__init__.py b/rsl_rl/build/lib/rsl_rl/utils/__init__.py new file mode 100644 index 0000000..f5781f1 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/utils/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Helper functions.""" + +from .utils import * +from .motion_loader_g1 import G1_AMPLoader + +__all__ = [ + "G1_AMPLoader", +] \ No newline at end of file diff --git a/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py b/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py new file mode 100644 index 0000000..677c630 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py @@ -0,0 +1,388 @@ +import os +from os.path import join as pjoin +import glob +import json +import logging + +import torch +import numpy as np +from pybullet_utils import transformations + +from rsl_rl.utils import motion_util + +_EPS = np.finfo(float).eps * 4.0 +def quaternion_slerp(q0, q1, fraction, spin=0, shortestpath=True): + """Batch quaternion spherical linear interpolation.""" + + out = torch.zeros_like(q0) + + zero_mask = torch.isclose(fraction, torch.zeros_like(fraction)).squeeze() + ones_mask = torch.isclose(fraction, torch.ones_like(fraction)).squeeze() + out[zero_mask] = q0[zero_mask] + out[ones_mask] = q1[ones_mask] + + d = torch.sum(q0 * q1, dim=-1, keepdim=True) + dist_mask = (torch.abs(torch.abs(d) - 1.0) < _EPS).squeeze() + out[dist_mask] = q0[dist_mask] + + if shortestpath: + d_old = torch.clone(d) + d = torch.where(d_old < 0, -d, d) + q1 = torch.where(d_old < 0, -q1, q1) + + angle = torch.acos(d) + spin * torch.pi + angle_mask = (torch.abs(angle) < _EPS).squeeze() + out[angle_mask] = q0[angle_mask] + + final_mask = torch.logical_or(zero_mask, ones_mask) + final_mask = torch.logical_or(final_mask, dist_mask) + final_mask = torch.logical_or(final_mask, angle_mask) + final_mask = torch.logical_not(final_mask) + + isin = 1.0 / angle + q0 *= torch.sin((1.0 - fraction) * angle) * isin + q1 *= torch.sin(fraction * angle) * isin + q0 += q1 + out[final_mask] = q0[final_mask] + return out + + +class G1_AMPLoader: + + def __init__( + self, + device, + time_between_frames, + motion_files, + preload_transitions=False, + num_preload_transitions=1000000, + num_frames=5, + ): + """Expert dataset provides AMP observations from Dog mocap dataset. + + time_between_frames: Amount of time in seconds between transition. + """ + self.device = device + self.time_between_frames = time_between_frames + self.num_frames = num_frames + + # Values to store for each trajectory. + self.trajectories = [] + self.trajectories_full = [] + self.trajectory_names = [] + self.trajectory_idxs = [] + self.trajectory_lens = [] # Traj length in seconds. + self.trajectory_weights = [] + self.trajectory_frame_durations = [] + self.trajectory_num_frames = [] + self.motion_dir = motion_files + # import ipdb; ipdb.set_trace() + for i, motion_file in enumerate(os.listdir(motion_files)): + self.trajectory_names.append(motion_file) + motion_path = pjoin(motion_files, motion_file) + motion_data = np.load(motion_path, allow_pickle=True) + motion_data_processed = np.zeros((motion_data.shape[0],36)) + + for f_i in range(motion_data.shape[0]): + motion_data_processed[f_i, :3] = motion_data[f_i, :3] # base pos + motion_data_processed[f_i, 3:7] = motion_data[f_i, 3:7] # base quat (wxyz) + motion_data_processed[f_i, 7:35] = motion_data[f_i, 7:35] # base vel + ''' + NOTE The order of motion_data_processed is + base pos 0:3, + base quat 3:7, wxyz + dof pos 7:36, (mujoco joint order) + ''' + self.trajectories.append(torch.tensor( + motion_data_processed[:, 7:], + dtype=torch.float32, + device=self.device + )) + + self.trajectories_full.append(torch.tensor( + motion_data_processed, + dtype=torch.float32, + device=self.device + )) + + self.trajectory_idxs.append(i) + self.trajectory_weights.append(1 / len(os.listdir(motion_files))) + frame_duration = 1 / 50 + + self.trajectory_frame_durations.append(frame_duration) + traj_len = (motion_data_processed.shape[0] - 1) * frame_duration # seconds + self.trajectory_lens.append(traj_len) + self.trajectory_num_frames.append(float(motion_data_processed.shape[0])) + print(f"Loaded {traj_len}s. motion from {motion_file}.") + + # Trajectory weights are used to sample some trajectories more than others. + self.trajectory_weights = np.array(self.trajectory_weights) / np.sum(self.trajectory_weights) + self.trajectory_frame_durations = np.array(self.trajectory_frame_durations) + self.trajectory_lens = np.array(self.trajectory_lens) + self.trajectory_num_frames = np.array(self.trajectory_num_frames) + + # Preload transitions. + self.preload_transitions = preload_transitions + if self.preload_transitions: + print(f'Preloading {num_preload_transitions} transitions') + + traj_idxs = self.weighted_traj_idx_sample_batch(num_preload_transitions) + times = self.traj_time_sample_batch(traj_idxs) + self.preloaded_s_prior = self.get_full_frame_at_time_batch(traj_idxs, times - self.time_between_frames) + self.preloaded_s = self.get_full_frame_at_time_batch(traj_idxs, times) + self.preloaded_s_next = self.get_full_frame_at_time_batch(traj_idxs, times + self.time_between_frames) + print(f'Finished preloading') + + # 预加载多帧数据 + self.preloaded_frames = [] + for i in range(self.num_frames): + frame_time = times + (i - (self.num_frames - 2)) * self.time_between_frames + full_frame = self.get_full_frame_at_time_batch(traj_idxs, frame_time) + # 预处理：提前提取并连接需要的列（7:26 和 29:33），避免每次生成时重复切片 + processed_frame = torch.cat([ + full_frame[:, 7:26], + full_frame[:, 29:33] + ], dim=-1) + self.preloaded_frames.append(processed_frame) + print(f'Finished preloading multiple frames') + + self.all_trajectories_full = torch.vstack(self.trajectories_full) + + def weighted_traj_idx_sample(self): + """Get traj idx via weighted sampling.""" + return np.random.choice( + self.trajectory_idxs, p=self.trajectory_weights) + + def weighted_traj_idx_sample_batch(self, size): + """Batch sample traj idxs.""" + return np.random.choice( + self.trajectory_idxs, size=size, p=self.trajectory_weights, + replace=True) + + def traj_time_sample(self, traj_idx): + """Sample random time for traj.""" + subst = self.time_between_frames + self.trajectory_frame_durations[traj_idx] + return max( + 0, (self.trajectory_lens[traj_idx] * np.random.uniform() - subst)) + + def traj_time_sample_batch(self, traj_idxs): + """Sample random time for multiple trajectories.""" + subst = self.time_between_frames + self.trajectory_frame_durations[traj_idxs] + time_samples = self.trajectory_lens[traj_idxs] * np.random.uniform(size=len(traj_idxs)) - subst + return np.maximum(np.zeros_like(time_samples), time_samples) + + def slerp(self, val0, val1, blend): + return (1.0 - blend) * val0 + blend * val1 + + def get_trajectory(self, traj_idx): + """Returns trajectory of AMP observations.""" + return self.trajectories_full[traj_idx] + + def get_frame_at_time(self, traj_idx, time): + """Returns frame for the given trajectory at the specified time.""" + p = float(time) / self.trajectory_lens[traj_idx] + n = self.trajectories[traj_idx].shape[0] + idx_low, idx_high = int(np.floor(p * n)), int(np.ceil(p * n)) + frame_start = self.trajectories[traj_idx][idx_low] + frame_end = self.trajectories[traj_idx][idx_high] + blend = p * n - idx_low + return self.slerp(frame_start, frame_end, blend) + + def get_frame_at_time_batch(self, traj_idxs, times): + """Returns frame for the given trajectory at the specified time.""" + p = times / self.trajectory_lens[traj_idxs] + n = self.trajectory_num_frames[traj_idxs] + idx_low, idx_high = np.floor(p * n).astype(np.int32), np.ceil(p * n).astype(np.int32) + all_frame_starts = torch.zeros(len(traj_idxs), self.observation_dim, device=self.device) + all_frame_ends = torch.zeros(len(traj_idxs), self.observation_dim, device=self.device) + for traj_idx in set(traj_idxs): + trajectory = self.trajectories[traj_idx] + traj_mask = traj_idxs == traj_idx + all_frame_starts[traj_mask] = trajectory[idx_low[traj_mask]] + all_frame_ends[traj_mask] = trajectory[idx_high[traj_mask]] + blend = torch.tensor(p * n - idx_low, device=self.device, dtype=torch.float32).unsqueeze(-1) + return self.slerp(all_frame_starts, all_frame_ends, blend) + + def get_full_frame_at_time(self, traj_idx, time): + """Returns full frame for the given trajectory at the specified time.""" + p = float(time) / self.trajectory_lens[traj_idx] + n = self.trajectories_full[traj_idx].shape[0] + idx_low, idx_high = int(np.floor(p * n)), int(np.ceil(p * n)) + frame_start = self.trajectories_full[traj_idx][idx_low] + frame_end = self.trajectories_full[traj_idx][idx_high] + blend = p * n - idx_low + print(idx_low, idx_high) + return self.blend_frame_pose(frame_start, frame_end, blend) + + def get_full_frame_at_time_batch(self, traj_idxs, times): + p = times / self.trajectory_lens[traj_idxs] + n = self.trajectory_num_frames[traj_idxs] + idx_low, idx_high = np.floor(p * n).astype(np.int32), np.ceil(p * n).astype(np.int32) + all_frame_pos_starts = torch.zeros(len(traj_idxs), 3, device=self.device) + all_frame_pos_ends = torch.zeros(len(traj_idxs), 3, device=self.device) + all_frame_rot_starts = torch.zeros(len(traj_idxs), 4, device=self.device) + all_frame_rot_ends = torch.zeros(len(traj_idxs), 4, device=self.device) + all_frame_amp_starts = torch.zeros(len(traj_idxs), 29, device=self.device) + all_frame_amp_ends = torch.zeros(len(traj_idxs), 29, device=self.device) + for traj_idx in set(traj_idxs): + trajectory = self.trajectories_full[traj_idx] + traj_mask = traj_idxs == traj_idx + all_frame_pos_starts[traj_mask] = G1_AMPLoader.get_root_pos_batch(trajectory[idx_low[traj_mask]]) + all_frame_pos_ends[traj_mask] = G1_AMPLoader.get_root_pos_batch(trajectory[idx_high[traj_mask]]) + all_frame_rot_starts[traj_mask] = G1_AMPLoader.get_root_rot_batch(trajectory[idx_low[traj_mask]]) + all_frame_rot_ends[traj_mask] = G1_AMPLoader.get_root_rot_batch(trajectory[idx_high[traj_mask]]) + all_frame_amp_starts[traj_mask] = trajectory[idx_low[traj_mask]][:, 7:36] # base vel3+ang3, dof vel23+ang23 + all_frame_amp_ends[traj_mask] = trajectory[idx_high[traj_mask]][:, 7:36] # base vel3+ang3, dof vel23+ang23 + blend = torch.tensor(p * n - idx_low, device=self.device, dtype=torch.float32).unsqueeze(-1) + pos_blend = self.slerp(all_frame_pos_starts, all_frame_pos_ends, blend) + rot_blend = quaternion_slerp(all_frame_rot_starts, all_frame_rot_ends, blend) + amp_blend = self.slerp(all_frame_amp_starts, all_frame_amp_ends, blend) + return torch.cat([pos_blend, rot_blend, amp_blend], dim=-1) + + def get_frame(self): + """Returns random frame.""" + traj_idx = self.weighted_traj_idx_sample() + sampled_time = self.traj_time_sample(traj_idx) + return self.get_frame_at_time(traj_idx, sampled_time) + + def get_full_frame(self): + """Returns random full frame.""" + traj_idx = self.weighted_traj_idx_sample() + sampled_time = self.traj_time_sample(traj_idx) + return self.get_full_frame_at_time(traj_idx, sampled_time) + + def get_full_frame_batch(self, num_frames): + if self.preload_transitions: + idxs = np.random.choice( + self.preloaded_s.shape[0], size=num_frames) + return self.preloaded_s[idxs] + else: + traj_idxs = self.weighted_traj_idx_sample_batch(num_frames) + times = self.traj_time_sample_batch(traj_idxs) + return self.get_full_frame_at_time_batch(traj_idxs, times) + + def blend_frame_pose(self, frame0, frame1, blend): + """Linearly interpolate between two frames, including orientation. + + Args: + frame0: First frame to be blended corresponds to (blend = 0). + frame1: Second frame to be blended corresponds to (blend = 1). + blend: Float between [0, 1], specifying the interpolation between + the two frames. + Returns: + An interpolation of the two frames. + """ + root_pos0, root_pos1 = G1_AMPLoader.get_root_pos(frame0), G1_AMPLoader.get_root_pos(frame1) + root_rot0, root_rot1 = G1_AMPLoader.get_root_rot(frame0), G1_AMPLoader.get_root_rot(frame1) + joints0, joints1 = G1_AMPLoader.get_joint_pose(frame0), G1_AMPLoader.get_joint_pose(frame1) + # tar_toe_pos_0, tar_toe_pos_1 = G1_AMPLoader.get_tar_toe_pos_local(frame0), G1_AMPLoader.get_tar_toe_pos_local(frame1) + linear_vel_0, linear_vel_1 = G1_AMPLoader.get_linear_vel(frame0), G1_AMPLoader.get_linear_vel(frame1) + angular_vel_0, angular_vel_1 = G1_AMPLoader.get_angular_vel(frame0), G1_AMPLoader.get_angular_vel(frame1) + joint_vel_0, joint_vel_1 = G1_AMPLoader.get_joint_vel(frame0), G1_AMPLoader.get_joint_vel(frame1) + + blend_root_pos = self.slerp(root_pos0, root_pos1, blend) + blend_root_rot = transformations.quaternion_slerp(root_rot0.cpu().numpy(), root_rot1.cpu().numpy(), blend) + blend_root_rot = torch.tensor(motion_util.standardize_quaternion(blend_root_rot),dtype=torch.float32, device=self.device) + blend_joints = self.slerp(joints0, joints1, blend) + # blend_tar_toe_pos = self.slerp(tar_toe_pos_0, tar_toe_pos_1, blend) + blend_linear_vel = self.slerp(linear_vel_0, linear_vel_1, blend) + blend_angular_vel = self.slerp(angular_vel_0, angular_vel_1, blend) + blend_joints_vel = self.slerp(joint_vel_0, joint_vel_1, blend) + + # return + # torch.cat([ + # blend_root_pos, blend_root_rot, blend_linear_vel, blend_angular_vel, blend_joints, blend_joints_vel]) + return torch.cat([blend_root_pos, blend_root_rot, blend_linear_vel, blend_angular_vel, blend_joints]) + + def feed_forward_generator_23dof_multi(self, num_mini_batch, mini_batch_size): + """Generates a batch of AMP transitions.""" + # import ipdb; ipdb.set_trace() + for _ in range(num_mini_batch): + if self.preload_transitions: + idxs = np.random.choice(self.preloaded_s.shape[0], size=mini_batch_size) + + frames = [] + for i in range(self.num_frames): + # 数据已在预加载时预处理，直接索引即可 + s = self.preloaded_frames[i][idxs] + frames.append(s) + else: + NotImplementedError('preload transition') + yield torch.stack(frames, dim=1) # [batch, num_frames, 16] + + + + + def quaternion_to_euler_array(self, quat): + # Ensure quaternion is in the correct format [x, y, z, w] + x, y, z, w =quat + + # Roll (x-axis rotation) + t0 = +2.0 * (w * x + y * z) + t1 = +1.0 - 2.0 * (x * x + y * y) + roll_x = np.arctan2(t0, t1) + + # Pitch (y-axis rotation) + t2 = +2.0 * (w * y - z * x) + t2 = np.clip(t2, -1.0, 1.0) + pitch_y = np.arcsin(t2) + + # Yaw (z-axis rotation) + t3 = +2.0 * (w * z + x * y) + t4 = +1.0 - 2.0 * (y * y + z * z) + yaw_z = np.arctan2(t3, t4) + + # Returns roll, pitch, yaw in a NumPy array in radians + return np.array([roll_x, pitch_y, yaw_z]) + + def euler_to_quaternion(self, root_rot): + roll, pitch, yaw = root_rot[0], root_rot[1], root_rot[2] + cy = np.cos(yaw * 0.5) + sy = np.sin(yaw * 0.5) + cp = np.cos(pitch * 0.5) + sp = np.sin(pitch * 0.5) + cr = np.cos(roll * 0.5) + sr = np.sin(roll * 0.5) + + qw = cy * cp * cr + sy * sp * sr + qx = cy * cp * sr - sy * sp * cr + qy = sy * cp * sr + cy * sp * cr + qz = sy * cp * cr - cy * sp * sr + + return np.array([qx, qy, qz, qw]) + + @property + def observation_dim(self): + """Size of AMP observations.""" + return self.trajectories[0].shape[1] + 1 + + @property + def num_motions(self): + return len(self.trajectory_names) + @staticmethod + def get_root_pos(pose): + return pose[0:3] + + @staticmethod + def get_root_pos_batch(poses): + return poses[:, 0:3] + + @staticmethod + def get_root_rot(pose): + return pose[3:7] + + @staticmethod + def get_root_rot_batch(poses): + return poses[:, 3:7] + + @staticmethod + def get_joint_pose_batch_12dof(poses): + return poses[:, 13:25] + + @staticmethod + def get_tar_toe_pos_local(pose): + return pose[G1_AMPLoader.TAR_TOE_POS_LOCAL_START_IDX:G1_AMPLoader.TAR_TOE_POS_LOCAL_END_IDX] + + @staticmethod + def get_tar_toe_pos_local_batch(poses): + return poses[:, G1_AMPLoader.TAR_TOE_POS_LOCAL_START_IDX:G1_AMPLoader.TAR_TOE_POS_LOCAL_END_IDX] diff --git a/rsl_rl/build/lib/rsl_rl/utils/motion_util.py b/rsl_rl/build/lib/rsl_rl/utils/motion_util.py new file mode 100644 index 0000000..3d49bfa --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/utils/motion_util.py @@ -0,0 +1,97 @@ +# coding=utf-8 +# Copyright 2020 The Google Research Authors. +# +# 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. + +"""Utility functions for processing motion clips.""" + +import os +import inspect +currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) +parentdir = os.path.dirname(os.path.dirname(currentdir)) +os.sys.path.insert(0, parentdir) + +import numpy as np + +from rsl_rl.utils import pose3d +# from pybullet_utils import transformations + + +def standardize_quaternion(q): + """Returns a quaternion where q.w >= 0 to remove redundancy due to q = -q. + + Args: + q: A quaternion to be standardized. + + Returns: + A quaternion with q.w >= 0. + + """ + if q[-1] < 0: + q = -q + return q + + +def normalize_rotation_angle(theta): + """Returns a rotation angle normalized between [-pi, pi]. + + Args: + theta: angle of rotation (radians). + + Returns: + An angle of rotation normalized between [-pi, pi]. + + """ + norm_theta = theta + if np.abs(norm_theta) > np.pi: + norm_theta = np.fmod(norm_theta, 2 * np.pi) + if norm_theta >= 0: + norm_theta += -2 * np.pi + else: + norm_theta += 2 * np.pi + + return norm_theta + + +def calc_heading(q): + """Returns the heading of a rotation q, specified as a quaternion. + + The heading represents the rotational component of q along the vertical + axis (z axis). + + Args: + q: A quaternion that the heading is to be computed from. + + Returns: + An angle representing the rotation about the z axis. + + """ + ref_dir = np.array([1, 0, 0]) + rot_dir = pose3d.QuaternionRotatePoint(ref_dir, q) + heading = np.arctan2(rot_dir[1], rot_dir[0]) + return heading + + +# def calc_heading_rot(q): +# """Return a quaternion representing the heading rotation of q along the vertical axis (z axis). + +# Args: +# q: A quaternion that the heading is to be computed from. + +# Returns: +# A quaternion representing the rotation about the z axis. + +# """ +# heading = calc_heading(q) +# q_heading = transformations.quaternion_about_axis(heading, [0, 0, 1]) +# return q_heading diff --git a/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py b/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py new file mode 100644 index 0000000..3796ec8 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py @@ -0,0 +1,94 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +from dataclasses import asdict +from torch.utils.tensorboard import SummaryWriter + +try: + import neptune +except ModuleNotFoundError: + raise ModuleNotFoundError("neptune-client is required to log to Neptune.") + + +class NeptuneLogger: + def __init__(self, project, token): + self.run = neptune.init_run(project=project, api_token=token) + + def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): + self.run["runner_cfg"] = runner_cfg + self.run["policy_cfg"] = policy_cfg + self.run["alg_cfg"] = alg_cfg + self.run["env_cfg"] = asdict(env_cfg) + + +class NeptuneSummaryWriter(SummaryWriter): + """Summary writer for Neptune.""" + + def __init__(self, log_dir: str, flush_secs: int, cfg): + super().__init__(log_dir, flush_secs) + + try: + project = cfg["neptune_project"] + except KeyError: + raise KeyError("Please specify neptune_project in the runner config, e.g. legged_gym.") + + try: + token = os.environ["NEPTUNE_API_TOKEN"] + except KeyError: + raise KeyError( + "Neptune api token not found. Please run or add to ~/.bashrc: export NEPTUNE_API_TOKEN=YOUR_API_TOKEN" + ) + + try: + entity = os.environ["NEPTUNE_USERNAME"] + except KeyError: + raise KeyError( + "Neptune username not found. Please run or add to ~/.bashrc: export NEPTUNE_USERNAME=YOUR_USERNAME" + ) + + neptune_project = entity + "/" + project + + self.neptune_logger = NeptuneLogger(neptune_project, token) + + self.name_map = { + "Train/mean_reward/time": "Train/mean_reward_time", + "Train/mean_episode_length/time": "Train/mean_episode_length_time", + } + + run_name = os.path.split(log_dir)[-1] + + self.neptune_logger.run["log_dir"].log(run_name) + + def _map_path(self, path): + if path in self.name_map: + return self.name_map[path] + else: + return path + + def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_style=False): + super().add_scalar( + tag, + scalar_value, + global_step=global_step, + walltime=walltime, + new_style=new_style, + ) + self.neptune_logger.run[self._map_path(tag)].log(scalar_value, step=global_step) + + def stop(self): + self.neptune_logger.run.stop() + + def log_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): + self.neptune_logger.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg) + + def save_model(self, model_path, iter): + self.neptune_logger.run["model/saved_model_" + str(iter)].upload(model_path) + + def save_file(self, path, iter=None): + name = path.rsplit("/", 1)[-1].split(".")[0] + self.neptune_logger.run["git_diff/" + name].upload(path) diff --git a/rsl_rl/build/lib/rsl_rl/utils/pose3d.py b/rsl_rl/build/lib/rsl_rl/utils/pose3d.py new file mode 100644 index 0000000..a4cba1e --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/utils/pose3d.py @@ -0,0 +1,283 @@ +# coding=utf-8 +# Copyright 2020 The Google Research Authors. +# +# 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. +"""Utilities for 3D pose conversion.""" +import math +import numpy as np + +# from pybullet_utils import transformations + +VECTOR3_0 = np.zeros(3, dtype=np.float64) +VECTOR3_1 = np.ones(3, dtype=np.float64) +VECTOR3_X = np.array([1, 0, 0], dtype=np.float64) +VECTOR3_Y = np.array([0, 1, 0], dtype=np.float64) +VECTOR3_Z = np.array([0, 0, 1], dtype=np.float64) + +# QUATERNION_IDENTITY is the multiplicative identity 1.0 + 0i + 0j + 0k. +# When interpreted as a rotation, it is the identity rotation. +QUATERNION_IDENTITY = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float64) + + +def Vector3RandomNormal(sigma, mu=VECTOR3_0): + """Returns a random 3D vector from a normal distribution. + + Each component is selected independently from a normal distribution. + + Args: + sigma: Scale (or stddev) of distribution for all variables. + mu: Mean of distribution for each variable. + + Returns: + A 3D vector in a numpy array. + """ + + random_v3 = np.random.normal(scale=sigma, size=3) + mu + return random_v3 + + +def Vector3RandomUniform(low=VECTOR3_0, high=VECTOR3_1): + """Returns a 3D vector selected uniformly from the input box. + + Args: + low: The min-value corner of the box. + high: The max-value corner of the box. + + Returns: + A 3D vector in a numpy array. + """ + + random_x = np.random.uniform(low=low[0], high=high[0]) + random_y = np.random.uniform(low=low[1], high=high[1]) + random_z = np.random.uniform(low=low[2], high=high[2]) + return np.array([random_x, random_y, random_z]) + + +def Vector3RandomUnit(): + """Returns a random 3D vector with unit length. + + Generates a 3D vector selected uniformly from the unit sphere. + + Returns: + A normalized 3D vector in a numpy array. + """ + longitude = np.random.uniform(low=-math.pi, high=math.pi) + sin_latitude = np.random.uniform(low=-1.0, high=1.0) + cos_latitude = math.sqrt(1.0 - sin_latitude * sin_latitude) + x = math.cos(longitude) * cos_latitude + y = math.sin(longitude) * cos_latitude + z = sin_latitude + return np.array([x, y, z], dtype=np.float64) + + +def QuaternionNormalize(q): + """Normalizes the quaternion to length 1. + + Divides the quaternion by its magnitude. If the magnitude is too + small, returns the quaternion identity value (1.0). + + Args: + q: A quaternion to be normalized. + + Raises: + ValueError: If input quaternion has length near zero. + + Returns: + A quaternion with magnitude 1 in a numpy array [x, y, z, w]. + + """ + q_norm = np.linalg.norm(q) + if np.isclose(q_norm, 0.0): + raise ValueError( + 'Quaternion may not be zero in QuaternionNormalize: |q| = %f, q = %s' % + (q_norm, q)) + return q / q_norm + + +def QuaternionFromAxisAngle(axis, angle): + """Returns a quaternion that generates the given axis-angle rotation. + + Returns the quaternion: sin(angle/2) * axis + cos(angle/2). + + Args: + axis: Axis of rotation, a 3D vector in a numpy array. + angle: The angle of rotation (radians). + + Raises: + ValueError: If input axis is not a normalizable 3D vector. + + Returns: + A unit quaternion in a numpy array. + + """ + if len(axis) != 3: + raise ValueError('Axis vector should have three components: %s' % axis) + axis_norm = np.linalg.norm(axis) + if np.isclose(axis_norm, 0.0): + raise ValueError('Axis vector may not have zero length: |v| = %f, v = %s' % + (axis_norm, axis)) + half_angle = angle * 0.5 + q = np.zeros(4, dtype=np.float64) + q[0:3] = axis + q[0:3] *= math.sin(half_angle) / axis_norm + q[3] = math.cos(half_angle) + return q + + +def QuaternionToAxisAngle(quat, default_axis=VECTOR3_Z, direction_axis=None): + """Calculates axis and angle of rotation performed by a quaternion. + + Calculates the axis and angle of the rotation performed by the quaternion. + The quaternion should have four values and be normalized. + + Args: + quat: Unit quaternion in a numpy array. + default_axis: 3D vector axis used if the rotation is near to zero. Without + this default, small rotations would result in an exception. It is + reasonable to use a default axis for tiny rotations, because zero angle + rotations about any axis are equivalent. + direction_axis: Used to disambiguate rotation directions. If the + direction_axis is specified, the axis of the rotation will be chosen such + that its inner product with the direction_axis is non-negative. + + Raises: + ValueError: If quat is not a normalized quaternion. + + Returns: + axis: Axis of rotation. + angle: Angle in radians. + """ + if len(quat) != 4: + raise ValueError( + 'Quaternion should have four components [x, y, z, w]: %s' % quat) + if not np.isclose(1.0, np.linalg.norm(quat)): + raise ValueError('Quaternion should have unit length: |q| = %f, q = %s' % + (np.linalg.norm(quat), quat)) + axis = quat[:3].copy() + axis_norm = np.linalg.norm(axis) + min_axis_norm = 1e-8 + if axis_norm < min_axis_norm: + axis = default_axis + if len(default_axis) != 3: + raise ValueError('Axis vector should have three components: %s' % axis) + if not np.isclose(np.linalg.norm(axis), 1.0): + raise ValueError('Axis vector should have unit length: |v| = %f, v = %s' % + (np.linalg.norm(axis), axis)) + else: + axis /= axis_norm + sin_half_angle = axis_norm + if direction_axis is not None and np.inner(axis, direction_axis) < 0: + sin_half_angle = -sin_half_angle + axis = -axis + cos_half_angle = quat[3] + half_angle = math.atan2(sin_half_angle, cos_half_angle) + angle = half_angle * 2 + return axis, angle + + +def QuaternionRandomRotation(max_angle=math.pi): + """Creates a random small rotation around a random axis. + + Generates a small rotation with the axis vector selected uniformly + from the unit sphere and an angle selected from a uniform + distribution over [0, max_angle]. + + If the max_angle is not specified, the rotation should be selected + uniformly over all possible rotation angles. + + Args: + max_angle: The maximum angle of rotation (radians). + + Returns: + A unit quaternion in a numpy array. + + """ + + angle = np.random.uniform(low=0, high=max_angle) + axis = Vector3RandomUnit() + return QuaternionFromAxisAngle(axis, angle) + + +# def QuaternionRotatePoint(point, quat): +# """Performs a rotation by quaternion. + +# Rotate the point by the quaternion using quaternion multiplication, +# (q * p * q^-1), without constructing the rotation matrix. + +# Args: +# point: The point to be rotated. +# quat: The rotation represented as a quaternion [x, y, z, w]. + +# Returns: +# A 3D vector in a numpy array. +# """ + +# q_point = np.array([point[0], point[1], point[2], 0.0]) +# quat_inverse = transformations.quaternion_inverse(quat) +# q_point_rotated = transformations.quaternion_multiply( +# transformations.quaternion_multiply(quat, q_point), quat_inverse) +# return q_point_rotated[:3] + + +def IsRotationMatrix(m): + """Returns true if the 3x3 submatrix represents a rotation. + + Args: + m: A transformation matrix. + + Raises: + ValueError: If input is not a matrix of size at least 3x3. + + Returns: + True if the 3x3 submatrix is a rotation (orthogonal). + """ + if len(m.shape) != 2 or m.shape[0] < 3 or m.shape[1] < 3: + raise ValueError('Matrix should be 3x3 or 4x4: %s\n %s' % (m.shape, m)) + rot = m[:3, :3] + eye = np.matmul(rot, np.transpose(rot)) + return np.isclose(eye, np.identity(3), atol=1e-4).all() + +# def ZAxisAlignedRobotPoseTool(robot_pose_tool): +# """Returns the current gripper pose rotated for alignment with the z-axis. + +# Args: +# robot_pose_tool: a pose3d.Pose3d() instance. + +# Returns: +# An instance of pose.Transform representing the current gripper pose +# rotated for alignment with the z-axis. +# """ +# # Align the current pose to the z-axis. +# robot_pose_tool.quaternion = transformations.quaternion_multiply( +# RotationBetween( +# robot_pose_tool.matrix4x4[0:3, 0:3].dot(np.array([0, 0, 1])), +# np.array([0.0, 0.0, -1.0])), robot_pose_tool.quaternion) +# return robot_pose_tool + +# def RotationBetween(a_translation_b, a_translation_c): +# """Computes the rotation from one vector to another. + +# The computed rotation has the property that: + +# a_translation_c = a_rotation_b_to_c * a_translation_b + +# Args: +# a_translation_b: vec3, vector to rotate from +# a_translation_c: vec3, vector to rotate to + +# Returns: +# a_rotation_b_to_c: new Orientation +# """ +# rotation = rotation3.Rotation3.rotation_between( +# a_translation_b, a_translation_c, err_msg='RotationBetween') +# return rotation.quaternion.xyzw diff --git a/rsl_rl/build/lib/rsl_rl/utils/utils.py b/rsl_rl/build/lib/rsl_rl/utils/utils.py new file mode 100644 index 0000000..3605622 --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/utils/utils.py @@ -0,0 +1,360 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import git +import importlib +import os +import pathlib +import torch +import warnings +from tensordict import TensorDict +from typing import Callable +import numpy as np +class RunningMeanStd: + def __init__(self, epsilon: float = 1e-4, shape: Tuple[int, ...] = ()): + """ + Calculates the running mean and std of a data stream + https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm + :param epsilon: helps with arithmetic issues + :param shape: the shape of the data stream's output + """ + self.mean = np.zeros(shape, np.float64) + self.var = np.ones(shape, np.float64) + self.count = epsilon + + def update(self, arr: np.ndarray) -> None: + batch_mean = np.mean(arr, axis=0) + batch_var = np.var(arr, axis=0) + batch_count = arr.shape[0] + self.update_from_moments(batch_mean, batch_var, batch_count) + + def update_from_moments(self, batch_mean: np.ndarray, batch_var: np.ndarray, batch_count: int) -> None: + delta = batch_mean - self.mean + tot_count = self.count + batch_count + + new_mean = self.mean + delta * batch_count / tot_count + m_a = self.var * self.count + m_b = batch_var * batch_count + m_2 = m_a + m_b + np.square(delta) * self.count * batch_count / (self.count + batch_count) + new_var = m_2 / (self.count + batch_count) + + new_count = batch_count + self.count + + self.mean = new_mean + self.var = new_var + self.count = new_count + + +class Normalizer(RunningMeanStd): + def __init__(self, input_dim, epsilon=1e-4, clip_obs=10.0): + super().__init__(shape=input_dim) + self.epsilon = epsilon + self.clip_obs = clip_obs + + def normalize(self, input): + return np.clip((input - self.mean) / np.sqrt(self.var + self.epsilon), -self.clip_obs, self.clip_obs) + + def normalize_torch(self, input, device): + mean_torch = torch.tensor(self.mean, device=device, dtype=torch.float32) + std_torch = torch.sqrt(torch.tensor(self.var + self.epsilon, device=device, dtype=torch.float32)) + return torch.clamp((input - mean_torch) / std_torch, -self.clip_obs, self.clip_obs) + + def update_normalizer(self, rollouts, expert_loader): + policy_data_generator = rollouts.feed_forward_generator_amp(None, mini_batch_size=expert_loader.batch_size) + expert_data_generator = expert_loader.dataset.feed_forward_generator_amp(expert_loader.batch_size) + + for expert_batch, policy_batch in zip(expert_data_generator, policy_data_generator): + self.update(torch.vstack(tuple(policy_batch) + tuple(expert_batch)).cpu().numpy()) + + +def resolve_nn_activation(act_name: str) -> torch.nn.Module: + """Resolves the activation function from the name. + + Args: + act_name: The name of the activation function. + + Returns: + The activation function. + + Raises: + ValueError: If the activation function is not found. + """ + act_dict = { + "elu": torch.nn.ELU(), + "selu": torch.nn.SELU(), + "relu": torch.nn.ReLU(), + "crelu": torch.nn.CELU(), + "lrelu": torch.nn.LeakyReLU(), + "tanh": torch.nn.Tanh(), + "sigmoid": torch.nn.Sigmoid(), + "softplus": torch.nn.Softplus(), + "gelu": torch.nn.GELU(), + "swish": torch.nn.SiLU(), + "mish": torch.nn.Mish(), + "identity": torch.nn.Identity(), + } + + act_name = act_name.lower() + if act_name in act_dict: + return act_dict[act_name] + else: + raise ValueError(f"Invalid activation function '{act_name}'. Valid activations are: {list(act_dict.keys())}") + + +def resolve_optimizer(optimizer_name: str) -> torch.optim.Optimizer: + """Resolves the optimizer from the name. + + Args: + optimizer_name: The name of the optimizer. + + Returns: + The optimizer. + + Raises: + ValueError: If the optimizer is not found. + """ + optimizer_dict = { + "adam": torch.optim.Adam, + "adamw": torch.optim.AdamW, + "sgd": torch.optim.SGD, + "rmsprop": torch.optim.RMSprop, + } + + optimizer_name = optimizer_name.lower() + if optimizer_name in optimizer_dict: + return optimizer_dict[optimizer_name] + else: + raise ValueError(f"Invalid optimizer '{optimizer_name}'. Valid optimizers are: {list(optimizer_dict.keys())}") + + +def split_and_pad_trajectories( + tensor: torch.Tensor | TensorDict, dones: torch.Tensor +) -> tuple[torch.Tensor | TensorDict, torch.Tensor]: + """Splits trajectories at done indices. Then concatenates them and pads with zeros up to the length of the longest + trajectory. Returns masks corresponding to valid parts of the trajectories. + + Example: + Input: [[a1, a2, a3, a4 | a5, a6], + [b1, b2 | b3, b4, b5 | b6]] + + Output:[[a1, a2, a3, a4], | [[True, True, True, True], + [a5, a6, 0, 0], | [True, True, False, False], + [b1, b2, 0, 0], | [True, True, False, False], + [b3, b4, b5, 0], | [True, True, True, False], + [b6, 0, 0, 0]] | [True, False, False, False]] + + Assumes that the input has the following order of dimensions: [time, number of envs, additional dimensions] + """ + + dones = dones.clone() + dones[-1] = 1 + # Permute the buffers to have order (num_envs, num_transitions_per_env, ...), for correct reshaping + flat_dones = dones.transpose(1, 0).reshape(-1, 1) + # Get length of trajectory by counting the number of successive not done elements + done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero()[:, 0])) + trajectory_lengths = done_indices[1:] - done_indices[:-1] + trajectory_lengths_list = trajectory_lengths.tolist() + # Extract the individual trajectories + if isinstance(tensor, TensorDict): + padded_trajectories = {} + for k, v in tensor.items(): + # split the tensor into trajectories + trajectories = torch.split(v.transpose(1, 0).flatten(0, 1), trajectory_lengths_list) + # add at least one full length trajectory + trajectories = trajectories + (torch.zeros(v.shape[0], *v.shape[2:], device=v.device),) + # pad the trajectories to the length of the longest trajectory + padded_trajectories[k] = torch.nn.utils.rnn.pad_sequence(trajectories) + # remove the added tensor + padded_trajectories[k] = padded_trajectories[k][:, :-1] + padded_trajectories = TensorDict( + padded_trajectories, batch_size=[tensor.batch_size[0], len(trajectory_lengths_list)] + ) + else: + # split the tensor into trajectories + trajectories = torch.split(tensor.transpose(1, 0).flatten(0, 1), trajectory_lengths_list) + # add at least one full length trajectory + trajectories = trajectories + (torch.zeros(tensor.shape[0], *tensor.shape[2:], device=tensor.device),) + # pad the trajectories to the length of the longest trajectory + padded_trajectories = torch.nn.utils.rnn.pad_sequence(trajectories) + # remove the added tensor + padded_trajectories = padded_trajectories[:, :-1] + # create masks for the valid parts of the trajectories + trajectory_masks = trajectory_lengths > torch.arange(0, tensor.shape[0], device=tensor.device).unsqueeze(1) + return padded_trajectories, trajectory_masks + + +def unpad_trajectories(trajectories, masks): + """Does the inverse operation of split_and_pad_trajectories()""" + # Need to transpose before and after the masking to have proper reshaping + return ( + trajectories.transpose(1, 0)[masks.transpose(1, 0)] + .view(-1, trajectories.shape[0], trajectories.shape[-1]) + .transpose(1, 0) + ) + + +def store_code_state(logdir, repositories) -> list: + git_log_dir = os.path.join(logdir, "git") + os.makedirs(git_log_dir, exist_ok=True) + file_paths = [] + for repository_file_path in repositories: + try: + repo = git.Repo(repository_file_path, search_parent_directories=True) + t = repo.head.commit.tree + except Exception: + print(f"Could not find git repository in {repository_file_path}. Skipping.") + # skip if not a git repository + continue + # get the name of the repository + repo_name = pathlib.Path(repo.working_dir).name + diff_file_name = os.path.join(git_log_dir, f"{repo_name}.diff") + # check if the diff file already exists + if os.path.isfile(diff_file_name): + continue + # write the diff file + print(f"Storing git diff for '{repo_name}' in: {diff_file_name}") + with open(diff_file_name, "x", encoding="utf-8") as f: + content = f"--- git status ---\n{repo.git.status()} \n\n\n--- git diff ---\n{repo.git.diff(t)}" + f.write(content) + # add the file path to the list of files to be uploaded + file_paths.append(diff_file_name) + return file_paths + + +def string_to_callable(name: str) -> Callable: + """Resolves the module and function names to return the function. + + Args: + name: The function name. The format should be 'module:attribute_name'. + + Raises: + ValueError: When the resolved attribute is not a function. + ValueError: When unable to resolve the attribute. + + Returns: + The function loaded from the module. + """ + try: + mod_name, attr_name = name.split(":") + mod = importlib.import_module(mod_name) + callable_object = getattr(mod, attr_name) + # check if attribute is callable + if callable(callable_object): + return callable_object + else: + raise ValueError(f"The imported object is not callable: '{name}'") + except AttributeError as e: + msg = ( + "We could not interpret the entry as a callable object. The format of input should be" + f" 'module:attribute_name'\nWhile processing input '{name}', received the error:\n {e}." + ) + raise ValueError(msg) + + +def resolve_obs_groups( + obs: TensorDict, obs_groups: dict[str, list[str]], default_sets: list[str] +) -> dict[str, list[str]]: + """Validates the observation configuration and defaults missing observation sets. + + The input is an observation dictionary `obs` containing observation groups and a configuration dictionary + `obs_groups` where the keys are the observation sets and the values are lists of observation groups. + + The configuration dictionary could for example look like: + { + "policy": ["group_1", "group_2"], + "critic": ["group_1", "group_3"] + } + + This means that the 'policy' observation set will contain the observations "group_1" and "group_2" and the + 'critic' observation set will contain the observations "group_1" and "group_3". This function will check that all + the observations in the 'policy' and 'critic' observation sets are present in the observation dictionary from the + environment. + + Additionally, if one of the `default_sets`, e.g. "critic", is not present in the configuration dictionary, + this function will: + + 1. Check if a group with the same name exists in the observations and assign this group to the observation set. + 2. If 1. fails, it will assign the observations from the 'policy' observation set to the default observation set. + + Args: + obs: Observations from the environment in the form of a dictionary. + obs_groups: Observation sets configuration. + default_sets: Reserved observation set names used by the algorithm (besides 'policy'). + If not provided in 'obs_groups', a default behavior gets triggered. + + Returns: + The resolved observation groups. + + Raises: + ValueError: If any observation set is an empty list. + ValueError: If any observation set contains an observation term that is not present in the observations. + """ + # check if policy observation set exists + if "policy" not in obs_groups.keys(): + if "policy" in obs: + obs_groups["policy"] = ["policy"] + warnings.warn( + "The observation configuration dictionary 'obs_groups' must contain the 'policy' key." + " As an observation group with the name 'policy' was found, this is assumed to be the observation set." + " Consider adding the 'policy' key to the 'obs_groups' dictionary for clarity." + " This behavior will be removed in a future version." + ) + else: + raise ValueError( + "The observation configuration dictionary 'obs_groups' must contain the 'policy' key." + f" Found keys: {list(obs_groups.keys())}" + ) + + # check all observation sets for valid observation groups + for set_name, groups in obs_groups.items(): + # check if the list is empty + if len(groups) == 0: + msg = f"The '{set_name}' key in the 'obs_groups' dictionary can not be an empty list." + if set_name in default_sets: + if set_name not in obs: + msg += " Consider removing the key to default to the observations used for the 'policy' set." + else: + msg += ( + f" Consider removing the key to default to the observation '{set_name}' from the environment." + ) + raise ValueError(msg) + # check groups exist inside the observations from the environment + for group in groups: + if group not in obs: + raise ValueError( + f"Observation '{group}' in observation set '{set_name}' not found in the observations from the" + f" environment. Available observations from the environment: {list(obs.keys())}" + ) + + # fill missing observation sets + for default_set_name in default_sets: + if default_set_name not in obs_groups.keys(): + if default_set_name in obs: + obs_groups[default_set_name] = [default_set_name] + warnings.warn( + f"The observation configuration dictionary 'obs_groups' must contain the '{default_set_name}' key." + f" As an observation group with the name '{default_set_name}' was found, this is assumed to be the" + f" observation set. Consider adding the '{default_set_name}' key to the 'obs_groups' dictionary for" + " clarity. This behavior will be removed in a future version." + ) + else: + obs_groups[default_set_name] = obs_groups["policy"].copy() + warnings.warn( + f"The observation configuration dictionary 'obs_groups' must contain the '{default_set_name}' key." + f" As the configuration for '{default_set_name}' is missing, the observations from the 'policy' set" + f" are used. Consider adding the '{default_set_name}' key to the 'obs_groups' dictionary for" + " clarity. This behavior will be removed in a future version." + ) + + # print the final parsed observation sets + print("-" * 80) + print("Resolved observation sets: ") + for set_name, groups in obs_groups.items(): + print("\t", set_name, ": ", groups) + print("-" * 80) + + return obs_groups diff --git a/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py b/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py new file mode 100644 index 0000000..243e82d --- /dev/null +++ b/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py @@ -0,0 +1,87 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +from dataclasses import asdict +from torch.utils.tensorboard import SummaryWriter + +try: + import wandb +except ModuleNotFoundError: + raise ModuleNotFoundError("Wandb is required to log to Weights and Biases.") + + +class WandbSummaryWriter(SummaryWriter): + """Summary writer for Weights and Biases.""" + + def __init__(self, log_dir: str, flush_secs: int, cfg): + super().__init__(log_dir, flush_secs) + + # Get the run name + run_name = os.path.split(log_dir)[-1] + + try: + project = cfg["wandb_project"] + except KeyError: + raise KeyError("Please specify wandb_project in the runner config, e.g. legged_gym.") + + try: + entity = os.environ["WANDB_USERNAME"] + except KeyError: + entity = None + + # Initialize wandb + wandb.init(project=project, entity=entity, name=run_name) + + # Add log directory to wandb + wandb.config.update({"log_dir": log_dir}) + + self.name_map = { + "Train/mean_reward/time": "Train/mean_reward_time", + "Train/mean_episode_length/time": "Train/mean_episode_length_time", + } + + def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): + wandb.config.update({"runner_cfg": runner_cfg}) + wandb.config.update({"policy_cfg": policy_cfg}) + wandb.config.update({"alg_cfg": alg_cfg}) + try: + wandb.config.update({"env_cfg": env_cfg.to_dict()}) + except Exception: + wandb.config.update({"env_cfg": asdict(env_cfg)}) + + def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_style=False): + super().add_scalar( + tag, + scalar_value, + global_step=global_step, + walltime=walltime, + new_style=new_style, + ) + wandb.log({self._map_path(tag): scalar_value}, step=global_step) + + def stop(self): + wandb.finish() + + def log_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): + self.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg) + + def save_model(self, model_path, iter): + wandb.save(model_path, base_path=os.path.dirname(model_path)) + + def save_file(self, path, iter=None): + wandb.save(path, base_path=os.path.dirname(path)) + + """ + Private methods. + """ + + def _map_path(self, path): + if path in self.name_map: + return self.name_map[path] + else: + return path diff --git a/rsl_rl/env/__init__.py b/rsl_rl/env/__init__.py new file mode 100644 index 0000000..ab7c056 --- /dev/null +++ b/rsl_rl/env/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Submodule defining the environment definitions.""" + +from .vec_env import VecEnv + +__all__ = ["VecEnv"] diff --git a/rsl_rl/env/__pycache__/__init__.cpython-312.pyc b/rsl_rl/env/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..044f56f Binary files /dev/null and b/rsl_rl/env/__pycache__/__init__.cpython-312.pyc differ diff --git a/rsl_rl/env/__pycache__/vec_env.cpython-312.pyc b/rsl_rl/env/__pycache__/vec_env.cpython-312.pyc new file mode 100644 index 0000000..9a9e657 Binary files /dev/null and b/rsl_rl/env/__pycache__/vec_env.cpython-312.pyc differ diff --git a/rsl_rl/env/vec_env.py b/rsl_rl/env/vec_env.py new file mode 100644 index 0000000..ae471cf --- /dev/null +++ b/rsl_rl/env/vec_env.py @@ -0,0 +1,113 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from abc import ABC, abstractmethod +from tensordict import TensorDict + + +class VecEnv(ABC): + """Abstract class for a vectorized environment. + + The vectorized environment is a collection of environments that are synchronized. This means that + the same type of action is applied to all environments and the same type of observation is returned from all + environments. + """ + + num_envs: int + """Number of environments.""" + + num_actions: int + """Number of actions.""" + + max_episode_length: int | torch.Tensor + + max_episode_length_s: float + """Maximum episode length. + + The maximum episode length can be a scalar or a tensor. If it is a scalar, it is the same for all environments. + If it is a tensor, it is the maximum episode length for each environment. This is useful for dynamic episode + lengths. + """ + + episode_length_buf: torch.Tensor + """Buffer for current episode lengths.""" + + device: torch.device | str + """Device to use.""" + + cfg: dict | object + """Configuration object.""" + + reset_env_ids: torch.Tensor | None = None + + contact_phase: torch.Tensor | None = None + """ + Operations. + """ + + @abstractmethod + def get_observations(self) -> TensorDict: + """Return the current observations. + + Returns: + observations (TensorDict): Observations from the environment. + """ + raise NotImplementedError + + @abstractmethod + def get_amp_observations(self) -> TensorDict: + """Return the current AMP observations. + + Returns: + observations (TensorDict): Observations from the environment. + """ + raise NotImplementedError + + @abstractmethod + def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]: + """Apply input action to the environment. + + Args: + actions (torch.Tensor): Input actions to apply. Shape: (num_envs, num_actions) + + Returns: + observations (TensorDict): Observations from the environment. + rewards (torch.Tensor): Rewards from the environment. Shape: (num_envs,) + dones (torch.Tensor): Done flags from the environment. Shape: (num_envs,) + extras (dict): Extra information from the environment. + + Observations: + + The observations TensorDict usually contains multiple observation groups. The `obs_groups` + dictionary of the runner configuration specifies which observation groups are used for which + purpose, i.e., it maps the available observation groups to observation sets. The observation sets + (keys of the `obs_groups` dictionary) currently used by rsl_rl are: + + - "policy": Specified observation groups are used as input to the actor/student network. + - "critic": Specified observation groups are used as input to the critic network. + - "teacher": Specified observation groups are used as input to the teacher network. + - "rnd_state": Specified observation groups are used as input to the RND network. + + Incomplete or incorrect configurations are handled in the `resolve_obs_groups()` function in + `rsl_rl/utils/utils.py`. + + Extras: + + The extras dictionary includes metrics such as the episode reward, episode length, etc. The following + dictionary keys are used by rsl_rl: + + - "time_outs" (torch.Tensor): Timeouts for the environments. These correspond to terminations that + happen due to time limits and not due to the environment reaching a terminal state. This is useful + for environments that have a fixed episode length. + + - "log" (dict[str, float | torch.Tensor]): Additional information for logging and debugging purposes. + The key should be a string and start with "/" for namespacing. The value can be a scalar or a + tensor. If it is a tensor, the mean of the tensor is used for logging. + """ + raise NotImplementedError + \ No newline at end of file diff --git a/rsl_rl/modules/__init__.py b/rsl_rl/modules/__init__.py new file mode 100644 index 0000000..9afe1a2 --- /dev/null +++ b/rsl_rl/modules/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Definitions for neural-network components for RL-agents.""" + +from .actor_critic import ActorCritic +from .actor_critic_recurrent import ActorCriticRecurrent +from .rnd import * +from .student_teacher import StudentTeacher +from .student_teacher_recurrent import StudentTeacherRecurrent +from .symmetry import * +from .discriminator_multi import DiscriminatorMulti +__all__ = [ + "ActorCritic", + "ActorCriticRecurrent", + "StudentTeacher", + "StudentTeacherRecurrent", + "DiscriminatorMulti", +] diff --git a/rsl_rl/modules/__pycache__/__init__.cpython-312.pyc b/rsl_rl/modules/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..69a4dbc Binary files /dev/null and b/rsl_rl/modules/__pycache__/__init__.cpython-312.pyc differ diff --git a/rsl_rl/modules/__pycache__/actor_critic.cpython-312.pyc b/rsl_rl/modules/__pycache__/actor_critic.cpython-312.pyc new file mode 100644 index 0000000..36400f3 Binary files /dev/null and b/rsl_rl/modules/__pycache__/actor_critic.cpython-312.pyc differ diff --git a/rsl_rl/modules/__pycache__/actor_critic_recurrent.cpython-312.pyc b/rsl_rl/modules/__pycache__/actor_critic_recurrent.cpython-312.pyc new file mode 100644 index 0000000..6c5e69e Binary files /dev/null and b/rsl_rl/modules/__pycache__/actor_critic_recurrent.cpython-312.pyc differ diff --git a/rsl_rl/modules/__pycache__/discriminator_multi.cpython-312.pyc b/rsl_rl/modules/__pycache__/discriminator_multi.cpython-312.pyc new file mode 100644 index 0000000..7d7b376 Binary files /dev/null and b/rsl_rl/modules/__pycache__/discriminator_multi.cpython-312.pyc differ diff --git a/rsl_rl/modules/__pycache__/rnd.cpython-312.pyc b/rsl_rl/modules/__pycache__/rnd.cpython-312.pyc new file mode 100644 index 0000000..ea7b356 Binary files /dev/null and b/rsl_rl/modules/__pycache__/rnd.cpython-312.pyc differ diff --git a/rsl_rl/modules/__pycache__/student_teacher.cpython-312.pyc b/rsl_rl/modules/__pycache__/student_teacher.cpython-312.pyc new file mode 100644 index 0000000..700cc03 Binary files /dev/null and b/rsl_rl/modules/__pycache__/student_teacher.cpython-312.pyc differ diff --git a/rsl_rl/modules/__pycache__/student_teacher_recurrent.cpython-312.pyc b/rsl_rl/modules/__pycache__/student_teacher_recurrent.cpython-312.pyc new file mode 100644 index 0000000..b707c3f Binary files /dev/null and b/rsl_rl/modules/__pycache__/student_teacher_recurrent.cpython-312.pyc differ diff --git a/rsl_rl/modules/__pycache__/symmetry.cpython-312.pyc b/rsl_rl/modules/__pycache__/symmetry.cpython-312.pyc new file mode 100644 index 0000000..db4e785 Binary files /dev/null and b/rsl_rl/modules/__pycache__/symmetry.cpython-312.pyc differ diff --git a/rsl_rl/modules/actor_critic.py b/rsl_rl/modules/actor_critic.py new file mode 100644 index 0000000..0efc36b --- /dev/null +++ b/rsl_rl/modules/actor_critic.py @@ -0,0 +1,195 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +from torch.distributions import Normal + +from rsl_rl.networks import MLP, EmpiricalNormalization + + +class ActorCritic(nn.Module): + is_recurrent = False + + def __init__( + self, + obs, + obs_groups, + num_actions, + actor_obs_normalization=False, + critic_obs_normalization=False, + actor_hidden_dims=[256, 256, 256], + critic_hidden_dims=[256, 256, 256], + activation="elu", + init_noise_std=1.0, + noise_std_type: str = "scalar", + state_dependent_std=False, + **kwargs, + ): + if kwargs: + print( + "ActorCritic.__init__ got unexpected arguments, which will be ignored: " + + str([key for key in kwargs.keys()]) + ) + super().__init__() + + # get the observation dimensions + self.obs_groups = obs_groups + num_actor_obs = 0 + for obs_group in obs_groups["policy"]: + assert len(obs[obs_group].shape) == 2, "The ActorCritic module only supports 1D observations." + num_actor_obs += obs[obs_group].shape[-1] + num_critic_obs = 0 + for obs_group in obs_groups["critic"]: + assert len(obs[obs_group].shape) == 2, "The ActorCritic module only supports 1D observations." + num_critic_obs += obs[obs_group].shape[-1] + + self.state_dependent_std = state_dependent_std + # actor + if self.state_dependent_std: + self.actor = MLP(num_actor_obs, [2, num_actions], actor_hidden_dims, activation) + else: + self.actor = MLP(num_actor_obs, num_actions, actor_hidden_dims, activation) + # actor observation normalization + self.actor_obs_normalization = actor_obs_normalization + if actor_obs_normalization: + self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs) + else: + self.actor_obs_normalizer = torch.nn.Identity() + print(f"Actor MLP: {self.actor}") + + # critic + self.critic = MLP(num_critic_obs, 1, critic_hidden_dims, activation) + # critic observation normalization + self.critic_obs_normalization = critic_obs_normalization + if critic_obs_normalization: + self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs) + else: + self.critic_obs_normalizer = torch.nn.Identity() + print(f"Critic MLP: {self.critic}") + + # Action noise + self.noise_std_type = noise_std_type + if self.state_dependent_std: + torch.nn.init.zeros_(self.actor[-2].weight[num_actions:]) + if self.noise_std_type == "scalar": + torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std) + elif self.noise_std_type == "log": + torch.nn.init.constant_( + self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7)) + ) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + else: + if self.noise_std_type == "scalar": + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + elif self.noise_std_type == "log": + self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + + # Action distribution (populated in update_distribution) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args(False) + + def reset(self, dones=None): + pass + + def forward(self): + raise NotImplementedError + + @property + def action_mean(self): + return self.distribution.mean + + @property + def action_std(self): + return self.distribution.stddev + + @property + def entropy(self): + return self.distribution.entropy().sum(dim=-1) + + def update_distribution(self, obs): + if self.state_dependent_std: + # compute mean and standard deviation + mean_and_std = self.actor(obs) + if self.noise_std_type == "scalar": + mean, std = torch.unbind(mean_and_std, dim=-2) + elif self.noise_std_type == "log": + mean, log_std = torch.unbind(mean_and_std, dim=-2) + std = torch.exp(log_std) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + else: + # compute mean + mean = self.actor(obs) + # compute standard deviation + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + # create distribution + self.distribution = Normal(mean, std) + + def act(self, obs, **kwargs): + obs = self.get_actor_obs(obs) + obs = self.actor_obs_normalizer(obs) + self.update_distribution(obs) + return self.distribution.sample() + + def act_inference(self, obs): + obs = self.get_actor_obs(obs) + obs = self.actor_obs_normalizer(obs) + return self.actor(obs) + + def evaluate(self, obs, **kwargs): + obs = self.get_critic_obs(obs) + obs = self.critic_obs_normalizer(obs) + return self.critic(obs) + + def get_actor_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["policy"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_critic_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["critic"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_actions_log_prob(self, actions): + return self.distribution.log_prob(actions).sum(dim=-1) + + def update_normalization(self, obs): + if self.actor_obs_normalization: + actor_obs = self.get_actor_obs(obs) + self.actor_obs_normalizer.update(actor_obs) + if self.critic_obs_normalization: + critic_obs = self.get_critic_obs(obs) + self.critic_obs_normalizer.update(critic_obs) + + def load_state_dict(self, state_dict, strict=True): + """Load the parameters of the actor-critic model. + + Args: + state_dict (dict): State dictionary of the model. + strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this + module's state_dict() function. + + Returns: + bool: Whether this training resumes a previous training. This flag is used by the `load()` function of + `OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation). + """ + + super().load_state_dict(state_dict, strict=strict) + return True # training resumes diff --git a/rsl_rl/modules/actor_critic_recurrent.py b/rsl_rl/modules/actor_critic_recurrent.py new file mode 100644 index 0000000..bba46ca --- /dev/null +++ b/rsl_rl/modules/actor_critic_recurrent.py @@ -0,0 +1,218 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +import warnings +from torch.distributions import Normal + +from rsl_rl.networks import MLP, EmpiricalNormalization, Memory + + +class ActorCriticRecurrent(nn.Module): + is_recurrent = True + + def __init__( + self, + obs, + obs_groups, + num_actions, + actor_obs_normalization=False, + critic_obs_normalization=False, + actor_hidden_dims=[256, 256, 256], + critic_hidden_dims=[256, 256, 256], + activation="elu", + init_noise_std=1.0, + noise_std_type: str = "scalar", + state_dependent_std=False, + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=1, + **kwargs, + ): + if "rnn_hidden_size" in kwargs: + warnings.warn( + "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. " + "Please use `rnn_hidden_dim` instead.", + DeprecationWarning, + ) + if rnn_hidden_dim == 256: # Only override if the new argument is at its default + rnn_hidden_dim = kwargs.pop("rnn_hidden_size") + if kwargs: + print( + "ActorCriticRecurrent.__init__ got unexpected arguments, which will be ignored: " + str(kwargs.keys()), + ) + super().__init__() + + # get the observation dimensions + self.obs_groups = obs_groups + num_actor_obs = 0 + for obs_group in obs_groups["policy"]: + assert len(obs[obs_group].shape) == 2, "The ActorCriticRecurrent module only supports 1D observations." + num_actor_obs += obs[obs_group].shape[-1] + num_critic_obs = 0 + for obs_group in obs_groups["critic"]: + assert len(obs[obs_group].shape) == 2, "The ActorCriticRecurrent module only supports 1D observations." + num_critic_obs += obs[obs_group].shape[-1] + + self.state_dependent_std = state_dependent_std + # actor + self.memory_a = Memory(num_actor_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) + if self.state_dependent_std: + self.actor = MLP(rnn_hidden_dim, [2, num_actions], actor_hidden_dims, activation) + else: + self.actor = MLP(rnn_hidden_dim, num_actions, actor_hidden_dims, activation) + + # actor observation normalization + self.actor_obs_normalization = actor_obs_normalization + if actor_obs_normalization: + self.actor_obs_normalizer = EmpiricalNormalization(num_actor_obs) + else: + self.actor_obs_normalizer = torch.nn.Identity() + print(f"Actor RNN: {self.memory_a}") + print(f"Actor MLP: {self.actor}") + + # critic + self.memory_c = Memory(num_critic_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) + self.critic = MLP(rnn_hidden_dim, 1, critic_hidden_dims, activation) + # critic observation normalization + self.critic_obs_normalization = critic_obs_normalization + if critic_obs_normalization: + self.critic_obs_normalizer = EmpiricalNormalization(num_critic_obs) + else: + self.critic_obs_normalizer = torch.nn.Identity() + print(f"Critic RNN: {self.memory_c}") + print(f"Critic MLP: {self.critic}") + + # Action noise + self.noise_std_type = noise_std_type + if self.state_dependent_std: + torch.nn.init.zeros_(self.actor[-2].weight[num_actions:]) + if self.noise_std_type == "scalar": + torch.nn.init.constant_(self.actor[-2].bias[num_actions:], init_noise_std) + elif self.noise_std_type == "log": + torch.nn.init.constant_( + self.actor[-2].bias[num_actions:], torch.log(torch.tensor(init_noise_std + 1e-7)) + ) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + else: + if self.noise_std_type == "scalar": + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + elif self.noise_std_type == "log": + self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + + # Action distribution (populated in update_distribution) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args(False) + + @property + def action_mean(self): + return self.distribution.mean + + @property + def action_std(self): + return self.distribution.stddev + + @property + def entropy(self): + return self.distribution.entropy().sum(dim=-1) + + def reset(self, dones=None): + self.memory_a.reset(dones) + self.memory_c.reset(dones) + + def forward(self): + raise NotImplementedError + + def update_distribution(self, obs): + if self.state_dependent_std: + # compute mean and standard deviation + mean_and_std = self.actor(obs) + if self.noise_std_type == "scalar": + mean, std = torch.unbind(mean_and_std, dim=-2) + elif self.noise_std_type == "log": + mean, log_std = torch.unbind(mean_and_std, dim=-2) + std = torch.exp(log_std) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + else: + # compute mean + mean = self.actor(obs) + # compute standard deviation + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + # create distribution + self.distribution = Normal(mean, std) + + def act(self, obs, masks=None, hidden_states=None): + obs = self.get_actor_obs(obs) + obs = self.actor_obs_normalizer(obs) + out_mem = self.memory_a(obs, masks, hidden_states).squeeze(0) + self.update_distribution(out_mem) + return self.distribution.sample() + + def act_inference(self, obs): + obs = self.get_actor_obs(obs) + obs = self.actor_obs_normalizer(obs) + out_mem = self.memory_a(obs).squeeze(0) + return self.actor(out_mem) + + def evaluate(self, obs, masks=None, hidden_states=None): + obs = self.get_critic_obs(obs) + obs = self.critic_obs_normalizer(obs) + out_mem = self.memory_c(obs, masks, hidden_states).squeeze(0) + return self.critic(out_mem) + + def get_actor_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["policy"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_critic_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["critic"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_actions_log_prob(self, actions): + return self.distribution.log_prob(actions).sum(dim=-1) + + def get_hidden_states(self): + return self.memory_a.hidden_states, self.memory_c.hidden_states + + def update_normalization(self, obs): + if self.actor_obs_normalization: + actor_obs = self.get_actor_obs(obs) + self.actor_obs_normalizer.update(actor_obs) + if self.critic_obs_normalization: + critic_obs = self.get_critic_obs(obs) + self.critic_obs_normalizer.update(critic_obs) + + def load_state_dict(self, state_dict, strict=True): + """Load the parameters of the actor-critic model. + + Args: + state_dict (dict): State dictionary of the model. + strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this + module's state_dict() function. + + Returns: + bool: Whether this training resumes a previous training. This flag is used by the `load()` function of + `OnPolicyRunner` to determine how to load further parameters (relevant for, e.g., distillation). + """ + + super().load_state_dict(state_dict, strict=strict) + return True diff --git a/rsl_rl/modules/discriminator_multi.py b/rsl_rl/modules/discriminator_multi.py new file mode 100644 index 0000000..7899d23 --- /dev/null +++ b/rsl_rl/modules/discriminator_multi.py @@ -0,0 +1,102 @@ +import torch +import torch.nn as nn +from torch import autograd +import torch.nn.utils.spectral_norm as spectral_norm + + +class DiscriminatorMulti(nn.Module): + def __init__( + self, state_dim, amp_reward_coef, hidden_layer_sizes, device, + num_frames=2, task_reward_lerp=0.0, use_lerp=True): + super(DiscriminatorMulti, self).__init__() + + self.device = device + self.state_dim = state_dim + self.use_lerp = use_lerp + self.num_frames = num_frames # 存储帧数参数 + + self.amp_reward_coef = amp_reward_coef + amp_layers = [] + + curr_in_dim = state_dim * num_frames + for hidden_dim in hidden_layer_sizes: + amp_layers.append(spectral_norm(nn.Linear(curr_in_dim, hidden_dim))) + amp_layers.append(nn.ReLU()) + curr_in_dim = hidden_dim + self.trunk = nn.Sequential(*amp_layers).to(device) + self.amp_linear = spectral_norm(nn.Linear(hidden_layer_sizes[-1], 1)).to(device) + + self.trunk.train() + self.amp_linear.train() + + self.task_reward_lerp = task_reward_lerp + + def forward(self, x): + h = self.trunk(x) + d = self.amp_linear(h) + return d + + def compute_grad_pen(self, + expert_states, # 改为接收多帧状态列表 + lambda_=10): + # 将多帧状态沿最后一个维度拼接 + expert_data = expert_states.flatten(1) + expert_data.requires_grad = True + + disc = self.amp_linear(self.trunk(expert_data)) + ones = torch.ones(disc.size(), device=disc.device) + grad = autograd.grad( + outputs=disc, inputs=expert_data, + grad_outputs=ones, create_graph=True, + retain_graph=True, only_inputs=True)[0] + + # Enforce that the grad norm approaches 0. + grad_pen = lambda_ * (grad.norm(2, dim=1) - 0).pow(2).mean() + return grad_pen + + + def get_disc_weights(self): + weights = [] + for m in self.trunk.modules(): + if isinstance(m, nn.Linear): + weights.append(torch.flatten(m.weight)) + + weights.append(torch.flatten(self.amp_linear.weight)) + return weights + + def get_disc_logit_weights(self): + return torch.flatten(self.amp_linear.weight) + + def predict_amp_reward( + self, states, # 改为接收多帧状态列表 + task_reward, normalizer=None): + """ + states: torch.Tensor, shape=(num_envs, num_frames, state_dim) + task_reward: torch.Tensor, shape=(num_envs, 1) + """ + # import ipdb; ipdb.set_trace() + with torch.no_grad(): + self.eval() + if normalizer is not None: + # 对每一帧状态进行归一化 + states = normalizer.normalize_torch(states, self.device) + + # 拼接多帧状态 + state_cat = states.flatten(1) + d = self.amp_linear(self.trunk(state_cat)) + disc_reward = self.amp_reward_coef * torch.clamp(1 - (1/4) * torch.square(d - 1), min=0) + + if self.use_lerp: + if self.task_reward_lerp > 0: + reward = self._lerp_reward(disc_reward, task_reward.unsqueeze(-1)) + self.train() + return reward.squeeze(), d, disc_reward.squeeze() * (1.0 - self.task_reward_lerp) + else: + disc_reward *= 0.02 + reward = task_reward.unsqueeze(-1) + disc_reward + self.train() + return reward.squeeze(), d, disc_reward.squeeze() + + def _lerp_reward(self, disc_r, task_r): + r = (1.0 - self.task_reward_lerp) * disc_r + self.task_reward_lerp * task_r + return r \ No newline at end of file diff --git a/rsl_rl/modules/rnd.py b/rsl_rl/modules/rnd.py new file mode 100644 index 0000000..8e65c43 --- /dev/null +++ b/rsl_rl/modules/rnd.py @@ -0,0 +1,209 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn + +from rsl_rl.networks import MLP, EmpiricalDiscountedVariationNormalization, EmpiricalNormalization + + +class RandomNetworkDistillation(nn.Module): + """Implementation of Random Network Distillation (RND) [1] + + References: + .. [1] Burda, Yuri, et al. "Exploration by random network distillation." arXiv preprint arXiv:1810.12894 (2018). + """ + + def __init__( + self, + num_states: int, + obs_groups: dict, + num_outputs: int, + predictor_hidden_dims: list[int], + target_hidden_dims: list[int], + activation: str = "elu", + weight: float = 0.0, + state_normalization: bool = False, + reward_normalization: bool = False, + device: str = "cpu", + weight_schedule: dict | None = None, + ): + """Initialize the RND module. + + - If :attr:`state_normalization` is True, then the input state is normalized using an Empirical Normalization layer. + - If :attr:`reward_normalization` is True, then the intrinsic reward is normalized using an Empirical Discounted + Variation Normalization layer. + + .. note:: + If the hidden dimensions are -1 in the predictor and target networks configuration, then the number of states + is used as the hidden dimension. + + Args: + num_states: Number of states/inputs to the predictor and target networks. + num_outputs: Number of outputs (embedding size) of the predictor and target networks. + predictor_hidden_dims: List of hidden dimensions of the predictor network. + target_hidden_dims: List of hidden dimensions of the target network. + activation: Activation function. Defaults to "elu". + weight: Scaling factor of the intrinsic reward. Defaults to 0.0. + state_normalization: Whether to normalize the input state. Defaults to False. + reward_normalization: Whether to normalize the intrinsic reward. Defaults to False. + device: Device to use. Defaults to "cpu". + weight_schedule: The type of schedule to use for the RND weight parameter. + Defaults to None, in which case the weight parameter is constant. + It is a dictionary with the following keys: + + - "mode": The type of schedule to use for the RND weight parameter. + - "constant": Constant weight schedule. + - "step": Step weight schedule. + - "linear": Linear weight schedule. + + For the "step" weight schedule, the following parameters are required: + + - "final_step": The step at which the weight parameter is set to the final value. + - "final_value": The final value of the weight parameter. + + For the "linear" weight schedule, the following parameters are required: + - "initial_step": The step at which the weight parameter is set to the initial value. + - "final_step": The step at which the weight parameter is set to the final value. + - "final_value": The final value of the weight parameter. + """ + # initialize parent class + super().__init__() + + # Store parameters + self.num_states = num_states + self.obs_groups = obs_groups + self.num_outputs = num_outputs + self.initial_weight = weight + self.device = device + self.state_normalization = state_normalization + self.reward_normalization = reward_normalization + + # Normalization of input gates + if state_normalization: + self.state_normalizer = EmpiricalNormalization(shape=[self.num_states], until=1.0e8).to(self.device) + else: + self.state_normalizer = torch.nn.Identity() + # Normalization of intrinsic reward + if reward_normalization: + self.reward_normalizer = EmpiricalDiscountedVariationNormalization(shape=[], until=1.0e8).to(self.device) + else: + self.reward_normalizer = torch.nn.Identity() + + # counter for the number of updates + self.update_counter = 0 + + # resolve weight schedule + if weight_schedule is not None: + self.weight_scheduler_params = weight_schedule + self.weight_scheduler = getattr(self, f"_{weight_schedule['mode']}_weight_schedule") + else: + self.weight_scheduler = None + # Create network architecture + self.predictor = MLP(num_states, num_outputs, predictor_hidden_dims, activation).to(self.device) + self.target = MLP(num_states, num_outputs, target_hidden_dims, activation).to(self.device) + + # make target network not trainable + self.target.eval() + + def get_intrinsic_reward(self, obs) -> torch.Tensor: + # Note: the counter is updated number of env steps per learning iteration + self.update_counter += 1 + # Extract the rnd state from the observation + rnd_state = self.get_rnd_state(obs) + rnd_state = self.state_normalizer(rnd_state) + # Obtain the embedding of the rnd state from the target and predictor networks + target_embedding = self.target(rnd_state).detach() + predictor_embedding = self.predictor(rnd_state).detach() + # Compute the intrinsic reward as the distance between the embeddings + intrinsic_reward = torch.linalg.norm(target_embedding - predictor_embedding, dim=1) + # Normalize intrinsic reward + intrinsic_reward = self.reward_normalizer(intrinsic_reward) + + # Check the weight schedule + if self.weight_scheduler is not None: + self.weight = self.weight_scheduler(step=self.update_counter, **self.weight_scheduler_params) + else: + self.weight = self.initial_weight + # Scale intrinsic reward + intrinsic_reward *= self.weight + + return intrinsic_reward + + def forward(self, *args, **kwargs): + raise RuntimeError("Forward method is not implemented. Use get_intrinsic_reward instead.") + + def train(self, mode: bool = True): + # sets module into training mode + self.predictor.train(mode) + if self.state_normalization: + self.state_normalizer.train(mode) + if self.reward_normalization: + self.reward_normalizer.train(mode) + return self + + def eval(self): + return self.train(False) + + def get_rnd_state(self, obs): + obs_list = [] + for obs_group in self.obs_groups["rnd_state"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def update_normalization(self, obs): + # Normalize the state + if self.state_normalization: + rnd_state = self.get_rnd_state(obs) + self.state_normalizer.update(rnd_state) + + """ + Different weight schedules. + """ + + def _constant_weight_schedule(self, step: int, **kwargs): + return self.initial_weight + + def _step_weight_schedule(self, step: int, final_step: int, final_value: float, **kwargs): + return self.initial_weight if step < final_step else final_value + + def _linear_weight_schedule(self, step: int, initial_step: int, final_step: int, final_value: float, **kwargs): + if step < initial_step: + return self.initial_weight + elif step > final_step: + return final_value + else: + return self.initial_weight + (final_value - self.initial_weight) * (step - initial_step) / ( + final_step - initial_step + ) + + +def resolve_rnd_config(alg_cfg, obs, obs_groups, env): + """Resolve the RND configuration. + + Args: + alg_cfg: The algorithm configuration dictionary. + obs: The observation dictionary. + obs_groups: The observation groups dictionary. + env: The environment. + + Returns: + The resolved algorithm configuration dictionary. + """ + # resolve dimension of rnd gated state + if "rnd_cfg" in alg_cfg and alg_cfg["rnd_cfg"] is not None: + # get dimension of rnd gated state + num_rnd_state = 0 + for obs_group in obs_groups["rnd_state"]: + assert len(obs[obs_group].shape) == 2, "The RND module only supports 1D observations." + num_rnd_state += obs[obs_group].shape[-1] + # add rnd gated state to config + alg_cfg["rnd_cfg"]["num_states"] = num_rnd_state + alg_cfg["rnd_cfg"]["obs_groups"] = obs_groups + # scale down the rnd weight with timestep + alg_cfg["rnd_cfg"]["weight"] *= env.unwrapped.step_dt + return alg_cfg diff --git a/rsl_rl/modules/student_teacher.py b/rsl_rl/modules/student_teacher.py new file mode 100644 index 0000000..6bf1380 --- /dev/null +++ b/rsl_rl/modules/student_teacher.py @@ -0,0 +1,206 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +from torch.distributions import Normal + +from rsl_rl.networks import MLP, EmpiricalNormalization + + +class StudentTeacher(nn.Module): + is_recurrent = False + + def __init__( + self, + obs, + obs_groups, + num_actions, + student_obs_normalization=False, + teacher_obs_normalization=False, + student_hidden_dims=[256, 256, 256], + teacher_hidden_dims=[256, 256, 256], + activation="elu", + init_noise_std=0.1, + noise_std_type: str = "scalar", + **kwargs, + ): + if kwargs: + print( + "StudentTeacher.__init__ got unexpected arguments, which will be ignored: " + + str([key for key in kwargs.keys()]) + ) + super().__init__() + + self.loaded_teacher = False # indicates if teacher has been loaded + + # get the observation dimensions + self.obs_groups = obs_groups + num_student_obs = 0 + for obs_group in obs_groups["policy"]: + assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." + num_student_obs += obs[obs_group].shape[-1] + num_teacher_obs = 0 + for obs_group in obs_groups["teacher"]: + assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." + num_teacher_obs += obs[obs_group].shape[-1] + + # student + self.student = MLP(num_student_obs, num_actions, student_hidden_dims, activation) + + # student observation normalization + self.student_obs_normalization = student_obs_normalization + if student_obs_normalization: + self.student_obs_normalizer = EmpiricalNormalization(num_student_obs) + else: + self.student_obs_normalizer = torch.nn.Identity() + + print(f"Student MLP: {self.student}") + + # teacher + self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation) + self.teacher.eval() + + # teacher observation normalization + self.teacher_obs_normalization = teacher_obs_normalization + if teacher_obs_normalization: + self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs) + else: + self.teacher_obs_normalizer = torch.nn.Identity() + + print(f"Teacher MLP: {self.teacher}") + + # action noise + self.noise_std_type = noise_std_type + if self.noise_std_type == "scalar": + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + elif self.noise_std_type == "log": + self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + + # action distribution (populated in update_distribution) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args(False) + + def reset(self, dones=None, hidden_states=None): + pass + + def forward(self): + raise NotImplementedError + + @property + def action_mean(self): + return self.distribution.mean + + @property + def action_std(self): + return self.distribution.stddev + + @property + def entropy(self): + return self.distribution.entropy().sum(dim=-1) + + def update_distribution(self, obs): + # compute mean + mean = self.student(obs) + # compute standard deviation + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + # create distribution + self.distribution = Normal(mean, std) + + def act(self, obs): + obs = self.get_student_obs(obs) + obs = self.student_obs_normalizer(obs) + self.update_distribution(obs) + return self.distribution.sample() + + def act_inference(self, obs): + obs = self.get_student_obs(obs) + obs = self.student_obs_normalizer(obs) + return self.student(obs) + + def evaluate(self, obs): + obs = self.get_teacher_obs(obs) + obs = self.teacher_obs_normalizer(obs) + with torch.no_grad(): + return self.teacher(obs) + + def get_student_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["policy"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_teacher_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["teacher"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_hidden_states(self): + return None + + def detach_hidden_states(self, dones=None): + pass + + def train(self, mode=True): + super().train(mode) + # make sure teacher is in eval mode + self.teacher.eval() + self.teacher_obs_normalizer.eval() + + def update_normalization(self, obs): + if self.student_obs_normalization: + student_obs = self.get_student_obs(obs) + self.student_obs_normalizer.update(student_obs) + + def load_state_dict(self, state_dict, strict=True): + """Load the parameters of the student and teacher networks. + + Args: + state_dict (dict): State dictionary of the model. + strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this + module's state_dict() function. + + Returns: + bool: Whether this training resumes a previous training. This flag is used by the `load()` function of + `OnPolicyRunner` to determine how to load further parameters. + """ + + # check if state_dict contains teacher and student or just teacher parameters + if any("actor" in key for key in state_dict.keys()): # loading parameters from rl training + # rename keys to match teacher and remove critic parameters + teacher_state_dict = {} + teacher_obs_normalizer_state_dict = {} + for key, value in state_dict.items(): + if "actor." in key: + teacher_state_dict[key.replace("actor.", "")] = value + if "actor_obs_normalizer." in key: + teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value + self.teacher.load_state_dict(teacher_state_dict, strict=strict) + self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict) + # set flag for successfully loading the parameters + self.loaded_teacher = True + self.teacher.eval() + self.teacher_obs_normalizer.eval() + return False # training does not resume + elif any("student" in key for key in state_dict.keys()): # loading parameters from distillation training + super().load_state_dict(state_dict, strict=strict) + # set flag for successfully loading the parameters + self.loaded_teacher = True + self.teacher.eval() + self.teacher_obs_normalizer.eval() + return True # training resumes + else: + raise ValueError("state_dict does not contain student or teacher parameters") diff --git a/rsl_rl/modules/student_teacher_recurrent.py b/rsl_rl/modules/student_teacher_recurrent.py new file mode 100644 index 0000000..964a2dc --- /dev/null +++ b/rsl_rl/modules/student_teacher_recurrent.py @@ -0,0 +1,249 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +import warnings +from torch.distributions import Normal + +from rsl_rl.networks import MLP, EmpiricalNormalization, Memory + + +class StudentTeacherRecurrent(nn.Module): + is_recurrent = True + + def __init__( + self, + obs, + obs_groups, + num_actions, + student_obs_normalization=False, + teacher_obs_normalization=False, + student_hidden_dims=[256, 256, 256], + teacher_hidden_dims=[256, 256, 256], + activation="elu", + init_noise_std=0.1, + noise_std_type: str = "scalar", + rnn_type="lstm", + rnn_hidden_dim=256, + rnn_num_layers=1, + teacher_recurrent=False, + **kwargs, + ): + if "rnn_hidden_size" in kwargs: + warnings.warn( + "The argument `rnn_hidden_size` is deprecated and will be removed in a future version. " + "Please use `rnn_hidden_dim` instead.", + DeprecationWarning, + ) + if rnn_hidden_dim == 256: # Only override if the new argument is at its default + rnn_hidden_dim = kwargs.pop("rnn_hidden_size") + if kwargs: + print( + "StudentTeacherRecurrent.__init__ got unexpected arguments, which will be ignored: " + + str(kwargs.keys()), + ) + super().__init__() + + self.loaded_teacher = False # indicates if teacher has been loaded + self.teacher_recurrent = teacher_recurrent # indicates if teacher is recurrent too + + # get the observation dimensions + self.obs_groups = obs_groups + num_student_obs = 0 + for obs_group in obs_groups["policy"]: + assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." + num_student_obs += obs[obs_group].shape[-1] + num_teacher_obs = 0 + for obs_group in obs_groups["teacher"]: + assert len(obs[obs_group].shape) == 2, "The StudentTeacher module only supports 1D observations." + num_teacher_obs += obs[obs_group].shape[-1] + + # student + self.memory_s = Memory(num_student_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim) + self.student = MLP(rnn_hidden_dim, num_actions, student_hidden_dims, activation) + + # student observation normalization + self.student_obs_normalization = student_obs_normalization + if student_obs_normalization: + self.student_obs_normalizer = EmpiricalNormalization(num_student_obs) + else: + self.student_obs_normalizer = torch.nn.Identity() + + print(f"Student RNN: {self.memory_s}") + print(f"Student MLP: {self.student}") + + # teacher + if self.teacher_recurrent: + self.memory_t = Memory( + num_teacher_obs, type=rnn_type, num_layers=rnn_num_layers, hidden_size=rnn_hidden_dim + ) + num_teacher_obs = rnn_hidden_dim + self.teacher = MLP(num_teacher_obs, num_actions, teacher_hidden_dims, activation) + + # teacher observation normalization + self.teacher_obs_normalization = teacher_obs_normalization + if teacher_obs_normalization: + self.teacher_obs_normalizer = EmpiricalNormalization(num_teacher_obs) + else: + self.teacher_obs_normalizer = torch.nn.Identity() + + if self.teacher_recurrent: + print(f"Teacher RNN: {self.memory_t}") + print(f"Teacher MLP: {self.teacher}") + + # action noise + self.noise_std_type = noise_std_type + if self.noise_std_type == "scalar": + self.std = nn.Parameter(init_noise_std * torch.ones(num_actions)) + elif self.noise_std_type == "log": + self.log_std = nn.Parameter(torch.log(init_noise_std * torch.ones(num_actions))) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + + # action distribution (populated in update_distribution) + self.distribution = None + # disable args validation for speedup + Normal.set_default_validate_args(False) + + def reset(self, dones=None, hidden_states=None): + if hidden_states is None: + hidden_states = (None, None) + self.memory_s.reset(dones, hidden_states[0]) + if self.teacher_recurrent: + self.memory_t.reset(dones, hidden_states[1]) + + def forward(self): + raise NotImplementedError + + @property + def action_mean(self): + return self.distribution.mean + + @property + def action_std(self): + return self.distribution.stddev + + @property + def entropy(self): + return self.distribution.entropy().sum(dim=-1) + + def update_distribution(self, obs): + # compute mean + mean = self.student(obs) + # compute standard deviation + if self.noise_std_type == "scalar": + std = self.std.expand_as(mean) + elif self.noise_std_type == "log": + std = torch.exp(self.log_std).expand_as(mean) + else: + raise ValueError(f"Unknown standard deviation type: {self.noise_std_type}. Should be 'scalar' or 'log'") + # create distribution + self.distribution = Normal(mean, std) + + def act(self, obs): + obs = self.get_student_obs(obs) + obs = self.student_obs_normalizer(obs) + out_mem = self.memory_s(obs).squeeze(0) + self.update_distribution(out_mem) + return self.distribution.sample() + + def act_inference(self, obs): + obs = self.get_student_obs(obs) + obs = self.student_obs_normalizer(obs) + out_mem = self.memory_s(obs).squeeze(0) + return self.student(out_mem) + + def evaluate(self, obs): + obs = self.get_teacher_obs(obs) + obs = self.teacher_obs_normalizer(obs) + with torch.no_grad(): + if self.teacher_recurrent: + self.memory_t.eval() + obs = self.memory_t(obs).squeeze(0) + return self.teacher(obs) + + def get_student_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["policy"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_teacher_obs(self, obs): + obs_list = [] + for obs_group in self.obs_groups["teacher"]: + obs_list.append(obs[obs_group]) + return torch.cat(obs_list, dim=-1) + + def get_hidden_states(self): + if self.teacher_recurrent: + return self.memory_s.hidden_states, self.memory_t.hidden_states + else: + return self.memory_s.hidden_states, None + + def detach_hidden_states(self, dones=None): + self.memory_s.detach_hidden_states(dones) + if self.teacher_recurrent: + self.memory_t.detach_hidden_states(dones) + + def train(self, mode=True): + super().train(mode) + # make sure teacher is in eval mode + self.teacher.eval() + self.teacher_obs_normalizer.eval() + + def update_normalization(self, obs): + if self.student_obs_normalization: + student_obs = self.get_student_obs(obs) + self.student_obs_normalizer.update(student_obs) + + def load_state_dict(self, state_dict, strict=True): + """Load the parameters of the student and teacher networks. + + Args: + state_dict (dict): State dictionary of the model. + strict (bool): Whether to strictly enforce that the keys in state_dict match the keys returned by this + module's state_dict() function. + + Returns: + bool: Whether this training resumes a previous training. This flag is used by the `load()` function of + `OnPolicyRunner` to determine how to load further parameters. + """ + + # check if state_dict contains teacher and student or just teacher parameters + if any("actor" in key for key in state_dict.keys()): # loading parameters from rl training + # rename keys to match teacher and remove critic parameters + teacher_state_dict = {} + teacher_obs_normalizer_state_dict = {} + for key, value in state_dict.items(): + if "actor." in key: + teacher_state_dict[key.replace("actor.", "")] = value + if "actor_obs_normalizer." in key: + teacher_obs_normalizer_state_dict[key.replace("actor_obs_normalizer.", "")] = value + self.teacher.load_state_dict(teacher_state_dict, strict=strict) + self.teacher_obs_normalizer.load_state_dict(teacher_obs_normalizer_state_dict, strict=strict) + # also load recurrent memory if teacher is recurrent + if self.teacher_recurrent: + memory_t_state_dict = {} + for key, value in state_dict.items(): + if "memory_a." in key: + memory_t_state_dict[key.replace("memory_a.", "")] = value + self.memory_t.load_state_dict(memory_t_state_dict, strict=strict) + # set flag for successfully loading the parameters + self.loaded_teacher = True + self.teacher.eval() + self.teacher_obs_normalizer.eval() + return False # training does not resume + elif any("student" in key for key in state_dict.keys()): # loading parameters from distillation training + super().load_state_dict(state_dict, strict=strict) + # set flag for successfully loading the parameters + self.loaded_teacher = True + self.teacher.eval() + self.teacher_obs_normalizer.eval() + return True # training resumes + else: + raise ValueError("state_dict does not contain student or teacher parameters") diff --git a/rsl_rl/modules/symmetry.py b/rsl_rl/modules/symmetry.py new file mode 100644 index 0000000..b017515 --- /dev/null +++ b/rsl_rl/modules/symmetry.py @@ -0,0 +1,24 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + + +def resolve_symmetry_config(alg_cfg, env): + """Resolve the symmetry configuration. + + Args: + alg_cfg: The algorithm configuration dictionary. + env: The environment. + + Returns: + The resolved algorithm configuration dictionary. + """ + + # if using symmetry then pass the environment config object + if "symmetry_cfg" in alg_cfg and alg_cfg["symmetry_cfg"] is not None: + # this is used by the symmetry function for handling different observation terms + alg_cfg["symmetry_cfg"]["_env"] = env + return alg_cfg diff --git a/rsl_rl/networks/__init__.py b/rsl_rl/networks/__init__.py new file mode 100644 index 0000000..c18f487 --- /dev/null +++ b/rsl_rl/networks/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Definitions for components of modules.""" + +from .memory import Memory +from .mlp import MLP +from .normalization import EmpiricalDiscountedVariationNormalization, EmpiricalNormalization diff --git a/rsl_rl/networks/__pycache__/__init__.cpython-312.pyc b/rsl_rl/networks/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..a79e7b4 Binary files /dev/null and b/rsl_rl/networks/__pycache__/__init__.cpython-312.pyc differ diff --git a/rsl_rl/networks/__pycache__/memory.cpython-312.pyc b/rsl_rl/networks/__pycache__/memory.cpython-312.pyc new file mode 100644 index 0000000..db9542b Binary files /dev/null and b/rsl_rl/networks/__pycache__/memory.cpython-312.pyc differ diff --git a/rsl_rl/networks/__pycache__/mlp.cpython-312.pyc b/rsl_rl/networks/__pycache__/mlp.cpython-312.pyc new file mode 100644 index 0000000..eb9fdfa Binary files /dev/null and b/rsl_rl/networks/__pycache__/mlp.cpython-312.pyc differ diff --git a/rsl_rl/networks/__pycache__/normalization.cpython-312.pyc b/rsl_rl/networks/__pycache__/normalization.cpython-312.pyc new file mode 100644 index 0000000..65433e7 Binary files /dev/null and b/rsl_rl/networks/__pycache__/normalization.cpython-312.pyc differ diff --git a/rsl_rl/networks/memory.py b/rsl_rl/networks/memory.py new file mode 100644 index 0000000..7577357 --- /dev/null +++ b/rsl_rl/networks/memory.py @@ -0,0 +1,70 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch.nn as nn + +from rsl_rl.utils import unpad_trajectories + + +class Memory(nn.Module): + """Memory module for recurrent networks. + + This module is used to store the hidden states of the policy. + Currently only supports GRU and LSTM. + """ + + def __init__(self, input_size, type="lstm", num_layers=1, hidden_size=256): + super().__init__() + # RNN + rnn_cls = nn.GRU if type.lower() == "gru" else nn.LSTM + self.rnn = rnn_cls(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers) + self.hidden_states = None + + def forward(self, input, masks=None, hidden_states=None): + batch_mode = masks is not None + if batch_mode: + # batch mode: needs saved hidden states + if hidden_states is None: + raise ValueError("Hidden states not passed to memory module during policy update") + out, _ = self.rnn(input, hidden_states) + out = unpad_trajectories(out, masks) + else: + # inference/distillation mode: uses hidden states of last step + out, self.hidden_states = self.rnn(input.unsqueeze(0), self.hidden_states) + return out + + def reset(self, dones=None, hidden_states=None): + if dones is None: # reset all hidden states + if hidden_states is None: + self.hidden_states = None + else: + self.hidden_states = hidden_states + elif self.hidden_states is not None: # reset hidden states of done environments + if hidden_states is None: + if isinstance(self.hidden_states, tuple): # tuple in case of LSTM + for hidden_state in self.hidden_states: + hidden_state[..., dones == 1, :] = 0.0 + else: + self.hidden_states[..., dones == 1, :] = 0.0 + else: + NotImplementedError( + "Resetting hidden states of done environments with custom hidden states is not implemented" + ) + + def detach_hidden_states(self, dones=None): + if self.hidden_states is not None: + if dones is None: # detach all hidden states + if isinstance(self.hidden_states, tuple): # tuple in case of LSTM + self.hidden_states = tuple(hidden_state.detach() for hidden_state in self.hidden_states) + else: + self.hidden_states = self.hidden_states.detach() + else: # detach hidden states of done environments + if isinstance(self.hidden_states, tuple): # tuple in case of LSTM + for hidden_state in self.hidden_states: + hidden_state[..., dones == 1, :] = hidden_state[..., dones == 1, :].detach() + else: + self.hidden_states[..., dones == 1, :] = self.hidden_states[..., dones == 1, :].detach() diff --git a/rsl_rl/networks/mlp.py b/rsl_rl/networks/mlp.py new file mode 100644 index 0000000..e91574e --- /dev/null +++ b/rsl_rl/networks/mlp.py @@ -0,0 +1,120 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +import torch.nn as nn +from functools import reduce + +from rsl_rl.utils import resolve_nn_activation + + +class MLP(nn.Sequential): + """Multi-layer perceptron. + + The MLP network is a sequence of linear layers and activation functions. The + last layer is a linear layer that outputs the desired dimension unless the + last activation function is specified. + + It provides additional conveniences: + + - If the hidden dimensions have a value of ``-1``, the dimension is inferred + from the input dimension. + - If the output dimension is a tuple, the output is reshaped to the desired + shape. + + """ + + def __init__( + self, + input_dim: int, + output_dim: int | tuple[int] | list[int], + hidden_dims: tuple[int] | list[int], + activation: str = "elu", + last_activation: str | None = None, + ): + """Initialize the MLP. + + Args: + input_dim: Dimension of the input. + output_dim: Dimension of the output. + hidden_dims: Dimensions of the hidden layers. A value of ``-1`` indicates + that the dimension should be inferred from the input dimension. + activation: Activation function. Defaults to "elu". + last_activation: Activation function of the last layer. Defaults to None, + in which case the last layer is linear. + """ + super().__init__() + + # resolve activation functions + activation_mod = resolve_nn_activation(activation) + last_activation_mod = resolve_nn_activation(last_activation) if last_activation is not None else None + # resolve number of hidden dims if they are -1 + hidden_dims_processed = [input_dim if dim == -1 else dim for dim in hidden_dims] + + # create layers sequentially + layers = [] + layers.append(nn.Linear(input_dim, hidden_dims_processed[0])) + layers.append(activation_mod) + + for layer_index in range(len(hidden_dims_processed) - 1): + layers.append(nn.Linear(hidden_dims_processed[layer_index], hidden_dims_processed[layer_index + 1])) + layers.append(activation_mod) + + # add last layer + if isinstance(output_dim, int): + layers.append(nn.Linear(hidden_dims_processed[-1], output_dim)) + else: + # compute the total output dimension + total_out_dim = reduce(lambda x, y: x * y, output_dim) + # add a layer to reshape the output to the desired shape + layers.append(nn.Linear(hidden_dims_processed[-1], total_out_dim)) + layers.append(nn.Unflatten(dim=-1, unflattened_size=output_dim)) + + # add last activation function if specified + if last_activation_mod is not None: + layers.append(last_activation_mod) + + # register the layers + for idx, layer in enumerate(layers): + self.add_module(f"{idx}", layer) + + def init_weights(self, scales: float | tuple[float]): + """Initialize the weights of the MLP. + + Args: + scales: Scale factor for the weights. + """ + + def get_scale(idx) -> float: + """Get the scale factor for the weights of the MLP. + + Args: + idx: Index of the layer. + """ + return scales[idx] if isinstance(scales, (list, tuple)) else scales + + # initialize the weights + for idx, module in enumerate(self): + if isinstance(module, nn.Linear): + nn.init.orthogonal_(module.weight, gain=get_scale(idx)) + nn.init.zeros_(module.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass of the MLP. + + Args: + x: Input tensor. + """ + for layer in self: + x = layer(x) + return x + + def reset(self, dones=None, hidden_states=None): + pass + + def detach_hidden_states(self, dones=None): + pass diff --git a/rsl_rl/networks/normalization.py b/rsl_rl/networks/normalization.py new file mode 100644 index 0000000..5fd9692 --- /dev/null +++ b/rsl_rl/networks/normalization.py @@ -0,0 +1,130 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Copyright (c) 2020 Preferred Networks, Inc. + +from __future__ import annotations + +import torch +from torch import nn + + +class EmpiricalNormalization(nn.Module): + """Normalize mean and variance of values based on empirical values.""" + + def __init__(self, shape, eps=1e-2, until=None): + """Initialize EmpiricalNormalization module. + + Args: + shape (int or tuple of int): Shape of input values except batch axis. + eps (float): Small value for stability. + until (int or None): If this arg is specified, the module learns input values until the sum of batch sizes + exceeds it. + + Note: The normalization parameters are computed over the whole batch, not for each environment separately. + """ + super().__init__() + self.eps = eps + self.until = until + self.register_buffer("_mean", torch.zeros(shape).unsqueeze(0)) + self.register_buffer("_var", torch.ones(shape).unsqueeze(0)) + self.register_buffer("_std", torch.ones(shape).unsqueeze(0)) + self.register_buffer("count", torch.tensor(0, dtype=torch.long)) + + @property + def mean(self): + return self._mean.squeeze(0).clone() + + @property + def std(self): + return self._std.squeeze(0).clone() + + def forward(self, x): + """Normalize mean and variance of values based on empirical values.""" + + return (x - self._mean) / (self._std + self.eps) + + @torch.jit.unused + def update(self, x): + """Learn input values without computing the output values of them""" + + if not self.training: + return + if self.until is not None and self.count >= self.until: + return + + count_x = x.shape[0] + self.count += count_x + rate = count_x / self.count + var_x = torch.var(x, dim=0, unbiased=False, keepdim=True) + mean_x = torch.mean(x, dim=0, keepdim=True) + delta_mean = mean_x - self._mean + self._mean += rate * delta_mean + self._var += rate * (var_x - self._var + delta_mean * (mean_x - self._mean)) + self._std = torch.sqrt(self._var) + + @torch.jit.unused + def inverse(self, y): + """De-normalize values based on empirical values.""" + + return y * (self._std + self.eps) + self._mean + + +class EmpiricalDiscountedVariationNormalization(nn.Module): + """Reward normalization from Pathak's large scale study on PPO. + + Reward normalization. Since the reward function is non-stationary, it is useful to normalize + the scale of the rewards so that the value function can learn quickly. We did this by dividing + the rewards by a running estimate of the standard deviation of the sum of discounted rewards. + """ + + def __init__(self, shape, eps=1e-2, gamma=0.99, until=None): + super().__init__() + + self.emp_norm = EmpiricalNormalization(shape, eps, until) + self.disc_avg = _DiscountedAverage(gamma) + + def forward(self, rew): + if self.training: + # update discounted rewards + avg = self.disc_avg.update(rew) + # update moments from discounted rewards + self.emp_norm.update(avg) + + # normalize rewards with the empirical std + if self.emp_norm._std > 0: + return rew / self.emp_norm._std + else: + return rew + + +""" +Helper class. +""" + + +class _DiscountedAverage: + r"""Discounted average of rewards. + + The discounted average is defined as: + + .. math:: + + \bar{R}_t = \gamma \bar{R}_{t-1} + r_t + + Args: + gamma (float): Discount factor. + """ + + def __init__(self, gamma): + self.avg = None + self.gamma = gamma + + def update(self, rew: torch.Tensor) -> torch.Tensor: + if self.avg is None: + self.avg = rew + else: + self.avg = self.avg * self.gamma + rew + return self.avg diff --git a/rsl_rl/pyproject.toml b/rsl_rl/pyproject.toml new file mode 100644 index 0000000..ae24da9 --- /dev/null +++ b/rsl_rl/pyproject.toml @@ -0,0 +1,23 @@ +[build-system] +requires = ["setuptools>=61", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "rsl-rl-lib" +version = "3.1.0" +description = "Fast and simple reinforcement learning algorithms (custom build)" +license = { text = "BSD-3-Clause" } +requires-python = ">=3.9" +dependencies = [ + "gitpython", + "numpy", + "onnx", + "tensordict", + "torch", + "torchvision", +] + +# 包在上一级目录（项目根下的 rsl_rl/），构建时从本目录运行故 where = ".." +[tool.setuptools.packages.find] +where = [".."] +include = ["rsl_rl*"] diff --git a/rsl_rl/runners/__init__.py b/rsl_rl/runners/__init__.py new file mode 100644 index 0000000..61f1682 --- /dev/null +++ b/rsl_rl/runners/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Implementation of runners for environment-agent interaction.""" + +from .on_policy_runner import OnPolicyRunner # isort:skip +from .distillation_runner import DistillationRunner +from .amp_on_policy_runner import AMPOnPolicyRunner + +__all__ = ["OnPolicyRunner", "DistillationRunner", "AMPOnPolicyRunner"] diff --git a/rsl_rl/runners/__pycache__/__init__.cpython-312.pyc b/rsl_rl/runners/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..4257c7a Binary files /dev/null and b/rsl_rl/runners/__pycache__/__init__.cpython-312.pyc differ diff --git a/rsl_rl/runners/__pycache__/amp_on_policy_runner.cpython-312.pyc b/rsl_rl/runners/__pycache__/amp_on_policy_runner.cpython-312.pyc new file mode 100644 index 0000000..da89059 Binary files /dev/null and b/rsl_rl/runners/__pycache__/amp_on_policy_runner.cpython-312.pyc differ diff --git a/rsl_rl/runners/__pycache__/distillation_runner.cpython-312.pyc b/rsl_rl/runners/__pycache__/distillation_runner.cpython-312.pyc new file mode 100644 index 0000000..2f967a4 Binary files /dev/null and b/rsl_rl/runners/__pycache__/distillation_runner.cpython-312.pyc differ diff --git a/rsl_rl/runners/__pycache__/on_policy_runner.cpython-312.pyc b/rsl_rl/runners/__pycache__/on_policy_runner.cpython-312.pyc new file mode 100644 index 0000000..76ef905 Binary files /dev/null and b/rsl_rl/runners/__pycache__/on_policy_runner.cpython-312.pyc differ diff --git a/rsl_rl/runners/amp_on_policy_runner.py b/rsl_rl/runners/amp_on_policy_runner.py new file mode 100644 index 0000000..c0b9b9e --- /dev/null +++ b/rsl_rl/runners/amp_on_policy_runner.py @@ -0,0 +1,521 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +import statistics +import time +import torch +import warnings +from collections import deque + +import rsl_rl +from rsl_rl.algorithms import AMP_PPO +from rsl_rl.env import VecEnv +from rsl_rl.modules import ActorCritic, ActorCriticRecurrent,DiscriminatorMulti, resolve_rnd_config, resolve_symmetry_config +from rsl_rl.utils import resolve_obs_groups, store_code_state, Normalizer, G1_AMPLoader + + +class AMPOnPolicyRunner: + """On-policy runner for training and evaluation of actor-critic methods.""" + + def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): + self.cfg = train_cfg + self.alg_cfg = train_cfg["algorithm"] + self.policy_cfg = train_cfg["policy"] + self.device = device + self.env = env + + # check if multi-gpu is enabled + self._configure_multi_gpu() + + # store training configuration + self.num_steps_per_env = self.cfg["num_steps_per_env"] + self.save_interval = self.cfg["save_interval"] + + # query observations from environment for algorithm construction + obs = self.env.get_observations() + default_sets = ["critic"] + if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None: + default_sets.append("rnd_state") + self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets) + + self.amp_data = G1_AMPLoader( + device, + time_between_frames=1/50.0, + preload_transitions=True, + num_preload_transitions=train_cfg["amp_num_preload_transitions"], + motion_files=train_cfg["amp_motion_files"], + num_frames=train_cfg['amp_num_frames'] + ) + + self.amp_observation_dim = self.amp_data.observation_dim if self.cfg["amp_num_obs"] == 0 else self.cfg["amp_num_obs"] + self.amp_num_frames = 0 if self.cfg["amp_num_frames"] == 0 else self.cfg["amp_num_frames"] + self.amp_normalizer = Normalizer(self.amp_observation_dim) + self.discriminator = DiscriminatorMulti( + self.amp_observation_dim, + train_cfg["amp_reward_coef"], + train_cfg["amp_discr_hidden_dims"], + device, + train_cfg["amp_num_frames"], + train_cfg["amp_task_reward_lerp"], + train_cfg['use_lerp'], + ).to(self.device) + + + # create the algorithm + self.alg = self._construct_algorithm(obs) + + # Decide whether to disable logging + # We only log from the process with rank 0 (main process) + self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 + + # Logging + self.log_dir = log_dir + self.writer = None + self.tot_timesteps = 0 + self.tot_time = 0 + self.current_learning_iteration = 0 + self.git_status_repos = [rsl_rl.__file__] + + def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 + # initialize writer + self._prepare_logging_writer() + + # randomize initial episode lengths (for exploration) + if init_at_random_ep_len: + self.env.episode_length_buf = torch.randint_like( + self.env.episode_length_buf, high=int(self.env.max_episode_length) + ) + + # start learning + obs = self.env.get_observations().to(self.device) + self.train_mode() # switch to train mode (for dropout for example) + amp_obs = self.env.get_amp_observations() + amp_obs = amp_obs.to(self.device) + if self.amp_num_frames != 0: + self.amp_obs_frames = torch.zeros(size=(self.env.num_envs, self.amp_num_frames, self.amp_observation_dim), device=self.device) + self.amp_obs_frames = torch.concat((self.amp_obs_frames[:, 1:], amp_obs.unsqueeze(1)), dim=1) + + # Book keeping + ep_infos = [] + rewbuffer = deque(maxlen=100) + lenbuffer = deque(maxlen=100) + step_discrewbuffer = deque(maxlen=100) + + cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_single_step_disc_rew = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + # create buffers for logging extrinsic and intrinsic rewards + if self.alg.rnd: + erewbuffer = deque(maxlen=100) + irewbuffer = deque(maxlen=100) + cur_ereward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_ireward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + # Ensure all parameters are in-synced + if self.is_distributed: + print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") + self.alg.broadcast_parameters() + + # Start training + start_iter = self.current_learning_iteration + tot_iter = start_iter + num_learning_iterations + for it in range(start_iter, tot_iter): + start = time.time() + # Rollout + with torch.inference_mode(): + for _ in range(self.num_steps_per_env): + # Sample actions + actions = self.alg.act(obs,amp_obs) + # Step the environment + obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) + # Move to device + obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) + + next_amp_obs = self.env.get_amp_observations() + next_amp_obs = next_amp_obs.to(self.device) + next_amp_obs_with_term = torch.clone(next_amp_obs) + + reset_env_ids = self.env.reset_env_ids + terminal_amp_states = self.env.get_amp_observations()[reset_env_ids] + next_amp_obs_with_term[reset_env_ids] = terminal_amp_states + self.amp_obs_frames = torch.concat((self.amp_obs_frames[:, 1:], next_amp_obs_with_term.unsqueeze(1)), dim=1) + + amp_reward = torch.zeros(self.env.num_envs, device=obs.device) + + mask = self.env.contact_phase[:, 0] == 1.0 + if mask.any(): + rewards[mask], logit, disc_reward = self.alg.discriminator.predict_amp_reward( + self.amp_obs_frames[mask], rewards[mask], normalizer=self.alg.amp_normalizer + ) + amp_reward[mask] += disc_reward + + # process the step + self.alg.process_env_step(obs, rewards, dones, extras, next_amp_obs_with_term, self.amp_obs_frames) + self.amp_obs_frames[reset_env_ids] = 0 + + amp_obs = torch.clone(next_amp_obs) + # Extract intrinsic rewards (only for logging) + intrinsic_rewards = self.alg.intrinsic_rewards if self.alg.rnd else None + # book keeping + if self.log_dir is not None: + if "episode" in extras: + ep_infos.append(extras["episode"]) + elif "log" in extras: + ep_infos.append(extras["log"]) + # Update rewards + if self.alg.rnd: + cur_ereward_sum += rewards + cur_ireward_sum += intrinsic_rewards # type: ignore + cur_reward_sum += rewards + intrinsic_rewards + else: + cur_reward_sum += rewards + # Update episode length + cur_episode_length += 1 + cur_single_step_disc_rew += amp_reward + # Clear data for completed episodes + # -- common + new_ids = (dones > 0).nonzero(as_tuple=False) + rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) + lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) + cur_reward_sum[new_ids] = 0 + cur_episode_length[new_ids] = 0 + to_extend_disc = (cur_single_step_disc_rew[new_ids] / self.env.max_episode_length_s)[:, 0].cpu().numpy() + step_discrewbuffer.extend(to_extend_disc.tolist()) + cur_single_step_disc_rew[new_ids] = 0 + # -- intrinsic and extrinsic rewards + if self.alg.rnd: + erewbuffer.extend(cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist()) + irewbuffer.extend(cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist()) + cur_ereward_sum[new_ids] = 0 + cur_ireward_sum[new_ids] = 0 + + stop = time.time() + collection_time = stop - start + start = stop + + # compute returns + self.alg.compute_returns(obs) + + # update policy + loss_dict = self.alg.update() + + stop = time.time() + learn_time = stop - start + self.current_learning_iteration = it + # log info + if self.log_dir is not None and not self.disable_logs: + # Log information + self.log(locals()) + # Save model + if it % self.save_interval == 0: + self.save(os.path.join(self.log_dir, f"model_{it}.pt")) + + # Clear episode infos + ep_infos.clear() + # Save code state + if it == start_iter and not self.disable_logs: + # obtain all the diff files + git_file_paths = store_code_state(self.log_dir, self.git_status_repos) + # if possible store them to wandb + if self.logger_type in ["wandb", "neptune"] and git_file_paths: + for path in git_file_paths: + self.writer.save_file(path) + + # Save the final model after training + if self.log_dir is not None and not self.disable_logs: + self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) + + def log(self, locs: dict, width: int = 80, pad: int = 35): + # Compute the collection size + collection_size = self.num_steps_per_env * self.env.num_envs * self.gpu_world_size + # Update total time-steps and time + self.tot_timesteps += collection_size + self.tot_time += locs["collection_time"] + locs["learn_time"] + iteration_time = locs["collection_time"] + locs["learn_time"] + + # -- Episode info + ep_string = "" + if locs["ep_infos"]: + for key in locs["ep_infos"][0]: + infotensor = torch.tensor([], device=self.device) + for ep_info in locs["ep_infos"]: + # handle scalar and zero dimensional tensor infos + if key not in ep_info: + continue + if not isinstance(ep_info[key], torch.Tensor): + ep_info[key] = torch.Tensor([ep_info[key]]) + if len(ep_info[key].shape) == 0: + ep_info[key] = ep_info[key].unsqueeze(0) + infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) + value = torch.mean(infotensor) + # log to logger and terminal + if "/" in key: + self.writer.add_scalar(key, value, locs["it"]) + ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + else: + self.writer.add_scalar("Episode/" + key, value, locs["it"]) + ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" + + mean_std = self.alg.policy.action_std.mean() + fps = int(collection_size / (locs["collection_time"] + locs["learn_time"])) + + # -- Losses + for key, value in locs["loss_dict"].items(): + self.writer.add_scalar(f"Loss/{key}", value, locs["it"]) + self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) + + # -- Policy + self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) + + # -- Performance + self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) + self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) + self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) + + # -- Training + if len(locs["rewbuffer"]) > 0: + # separate logging for intrinsic and extrinsic rewards + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(locs["erewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(locs["irewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/weight", self.alg.rnd.weight, locs["it"]) + # everything else + self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) + self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) + self.writer.add_scalar('Train/mean_step_disc_reward', statistics.mean(locs['step_discrewbuffer']), locs['it']) + if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging + self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) + self.writer.add_scalar( + "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time + ) + + str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " + + if len(locs["rewbuffer"]) > 0: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + f"""{'Step disc reward:':>{pad}} {statistics.mean(locs['step_discrewbuffer']):.2f}\n""" + ) + # -- Losses + for key, value in locs["loss_dict"].items(): + log_string += f"""{f'Mean {key} loss:':>{pad}} {value:.4f}\n""" + # -- Rewards + if hasattr(self.alg, "rnd") and self.alg.rnd: + log_string += ( + f"""{'Mean extrinsic reward:':>{pad}} {statistics.mean(locs['erewbuffer']):.2f}\n""" + f"""{'Mean intrinsic reward:':>{pad}} {statistics.mean(locs['irewbuffer']):.2f}\n""" + ) + log_string += f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" + # -- episode info + log_string += f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" + else: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + ) + for key, value in locs["loss_dict"].items(): + log_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + + log_string += ep_string + log_string += ( + f"""{'-' * width}\n""" + f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" + f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" + f"""{'Time elapsed:':>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n""" + f"""{'ETA:':>{pad}} {time.strftime( + "%H:%M:%S", + time.gmtime( + self.tot_time / (locs['it'] - locs['start_iter'] + 1) + * (locs['start_iter'] + locs['num_learning_iterations'] - locs['it']) + ) + )}\n""" + ) + print(log_string) + + def save(self, path: str, infos=None): + # -- Save model + saved_dict = { + "model_state_dict": self.alg.policy.state_dict(), + "optimizer_state_dict": self.alg.optimizer.state_dict(), + "iter": self.current_learning_iteration, + "infos": infos, + } + # -- Save RND model if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict() + saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict() + torch.save(saved_dict, path) + + # upload model to external logging service + if self.logger_type in ["neptune", "wandb"] and not self.disable_logs: + self.writer.save_model(path, self.current_learning_iteration) + + def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None): + loaded_dict = torch.load(path, weights_only=False, map_location=map_location) + # -- Load model + resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"]) + # -- Load RND model if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"]) + # -- load optimizer if used + if load_optimizer and resumed_training: + # -- algorithm optimizer + self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) + # -- RND optimizer if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"]) + # -- load current learning iteration + if resumed_training: + self.current_learning_iteration = loaded_dict["iter"] + return loaded_dict["infos"] + + def get_inference_policy(self, device=None): + self.eval_mode() # switch to evaluation mode (dropout for example) + if device is not None: + self.alg.policy.to(device) + return self.alg.policy.act_inference + + def train_mode(self): + # -- PPO + self.alg.policy.train() + self.alg.discriminator.train() + # -- RND + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.train() + + def eval_mode(self): + # -- PPO + self.alg.policy.eval() + self.alg.discriminator.eval() + # -- RND + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.eval() + + def add_git_repo_to_log(self, repo_file_path): + self.git_status_repos.append(repo_file_path) + + """ + Helper functions. + """ + + def _configure_multi_gpu(self): + """Configure multi-gpu training.""" + # check if distributed training is enabled + self.gpu_world_size = int(os.getenv("WORLD_SIZE", "1")) + self.is_distributed = self.gpu_world_size > 1 + + # if not distributed training, set local and global rank to 0 and return + if not self.is_distributed: + self.gpu_local_rank = 0 + self.gpu_global_rank = 0 + self.multi_gpu_cfg = None + return + + # get rank and world size + self.gpu_local_rank = int(os.getenv("LOCAL_RANK", "0")) + self.gpu_global_rank = int(os.getenv("RANK", "0")) + + # make a configuration dictionary + self.multi_gpu_cfg = { + "global_rank": self.gpu_global_rank, # rank of the main process + "local_rank": self.gpu_local_rank, # rank of the current process + "world_size": self.gpu_world_size, # total number of processes + } + + # check if user has device specified for local rank + if self.device != f"cuda:{self.gpu_local_rank}": + raise ValueError( + f"Device '{self.device}' does not match expected device for local rank '{self.gpu_local_rank}'." + ) + # validate multi-gpu configuration + if self.gpu_local_rank >= self.gpu_world_size: + raise ValueError( + f"Local rank '{self.gpu_local_rank}' is greater than or equal to world size '{self.gpu_world_size}'." + ) + if self.gpu_global_rank >= self.gpu_world_size: + raise ValueError( + f"Global rank '{self.gpu_global_rank}' is greater than or equal to world size '{self.gpu_world_size}'." + ) + + # initialize torch distributed + torch.distributed.init_process_group(backend="nccl", rank=self.gpu_global_rank, world_size=self.gpu_world_size) + # set device to the local rank + torch.cuda.set_device(self.gpu_local_rank) + + def _construct_algorithm(self, obs) -> AMP_PPO: + """Construct the actor-critic algorithm.""" + # resolve RND config + self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env) + + # resolve symmetry config + self.alg_cfg = resolve_symmetry_config(self.alg_cfg, self.env) + + # resolve deprecated normalization config + if self.cfg.get("empirical_normalization") is not None: + warnings.warn( + "The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and " + "`critic_obs_normalization` as part of the `policy` configuration instead.", + DeprecationWarning, + ) + if self.policy_cfg.get("actor_obs_normalization") is None: + self.policy_cfg["actor_obs_normalization"] = self.cfg["empirical_normalization"] + if self.policy_cfg.get("critic_obs_normalization") is None: + self.policy_cfg["critic_obs_normalization"] = self.cfg["empirical_normalization"] + + # initialize the actor-critic + actor_critic_class = eval(self.policy_cfg.pop("class_name")) + actor_critic: ActorCritic | ActorCriticRecurrent = actor_critic_class( + obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg + ).to(self.device) + + # initialize the algorithm + alg_class = eval(self.alg_cfg.pop("class_name")) + + alg: AMP_PPO = alg_class(actor_critic, self.discriminator, self.amp_data, self.amp_normalizer, self.amp_num_frames, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg) + + # initialize the storage + alg.init_storage( + "rl", + self.env.num_envs, + self.num_steps_per_env, + obs, + [self.env.num_actions], + ) + + return alg + + def _prepare_logging_writer(self): + """Prepares the logging writers.""" + if self.log_dir is not None and self.writer is None and not self.disable_logs: + # Launch either Tensorboard or Neptune & Tensorboard summary writer(s), default: Tensorboard. + self.logger_type = self.cfg.get("logger", "tensorboard") + self.logger_type = self.logger_type.lower() + + if self.logger_type == "neptune": + from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter + + self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "wandb": + from rsl_rl.utils.wandb_utils import WandbSummaryWriter + + self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "tensorboard": + from torch.utils.tensorboard import SummaryWriter + + self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) + else: + raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.") diff --git a/rsl_rl/runners/distillation_runner.py b/rsl_rl/runners/distillation_runner.py new file mode 100644 index 0000000..9cc6a8b --- /dev/null +++ b/rsl_rl/runners/distillation_runner.py @@ -0,0 +1,179 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +import time +import torch +from collections import deque + +import rsl_rl +from rsl_rl.algorithms import Distillation +from rsl_rl.env import VecEnv +from rsl_rl.modules import StudentTeacher, StudentTeacherRecurrent +from rsl_rl.runners import OnPolicyRunner +from rsl_rl.utils import resolve_obs_groups, store_code_state + + +class DistillationRunner(OnPolicyRunner): + """On-policy runner for training and evaluation of teacher-student training.""" + + def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): + self.cfg = train_cfg + self.alg_cfg = train_cfg["algorithm"] + self.policy_cfg = train_cfg["policy"] + self.device = device + self.env = env + + # check if multi-gpu is enabled + self._configure_multi_gpu() + + # store training configuration + self.num_steps_per_env = self.cfg["num_steps_per_env"] + self.save_interval = self.cfg["save_interval"] + + # query observations from environment for algorithm construction + obs = self.env.get_observations() + self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets=["teacher"]) + + # create the algorithm + self.alg = self._construct_algorithm(obs) + + # Decide whether to disable logging + # We only log from the process with rank 0 (main process) + self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 + + # Logging + self.log_dir = log_dir + self.writer = None + self.tot_timesteps = 0 + self.tot_time = 0 + self.current_learning_iteration = 0 + self.git_status_repos = [rsl_rl.__file__] + + def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 + # initialize writer + self._prepare_logging_writer() + # check if teacher is loaded + if not self.alg.policy.loaded_teacher: + raise ValueError("Teacher model parameters not loaded. Please load a teacher model to distill.") + + # randomize initial episode lengths (for exploration) + if init_at_random_ep_len: + self.env.episode_length_buf = torch.randint_like( + self.env.episode_length_buf, high=int(self.env.max_episode_length) + ) + + # start learning + obs = self.env.get_observations().to(self.device) + self.train_mode() # switch to train mode (for dropout for example) + + # Book keeping + ep_infos = [] + rewbuffer = deque(maxlen=100) + lenbuffer = deque(maxlen=100) + cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + # Ensure all parameters are in-synced + if self.is_distributed: + print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") + self.alg.broadcast_parameters() + + # Start training + start_iter = self.current_learning_iteration + tot_iter = start_iter + num_learning_iterations + for it in range(start_iter, tot_iter): + start = time.time() + # Rollout + with torch.inference_mode(): + for _ in range(self.num_steps_per_env): + # Sample actions + actions = self.alg.act(obs) + # Step the environment + obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) + # Move to device + obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) + # process the step + self.alg.process_env_step(obs, rewards, dones, extras) + # book keeping + if self.log_dir is not None: + if "episode" in extras: + ep_infos.append(extras["episode"]) + elif "log" in extras: + ep_infos.append(extras["log"]) + # Update rewards + cur_reward_sum += rewards + # Update episode length + cur_episode_length += 1 + # Clear data for completed episodes + new_ids = (dones > 0).nonzero(as_tuple=False) + rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) + lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) + cur_reward_sum[new_ids] = 0 + cur_episode_length[new_ids] = 0 + + stop = time.time() + collection_time = stop - start + start = stop + + # update policy + loss_dict = self.alg.update() + + stop = time.time() + learn_time = stop - start + self.current_learning_iteration = it + # log info + if self.log_dir is not None and not self.disable_logs: + # Log information + self.log(locals()) + # Save model + if it % self.save_interval == 0: + self.save(os.path.join(self.log_dir, f"model_{it}.pt")) + + # Clear episode infos + ep_infos.clear() + # Save code state + if it == start_iter and not self.disable_logs: + # obtain all the diff files + git_file_paths = store_code_state(self.log_dir, self.git_status_repos) + # if possible store them to wandb + if self.logger_type in ["wandb", "neptune"] and git_file_paths: + for path in git_file_paths: + self.writer.save_file(path) + + # Save the final model after training + if self.log_dir is not None and not self.disable_logs: + self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) + + """ + Helper methods. + """ + + def _construct_algorithm(self, obs) -> Distillation: + """Construct the distillation algorithm.""" + # initialize the actor-critic + student_teacher_class = eval(self.policy_cfg.pop("class_name")) + student_teacher: StudentTeacher | StudentTeacherRecurrent = student_teacher_class( + obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg + ).to(self.device) + + # initialize the algorithm + alg_class = eval(self.alg_cfg.pop("class_name")) + alg: Distillation = alg_class( + student_teacher, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg + ) + + # initialize the storage + alg.init_storage( + "distillation", + self.env.num_envs, + self.num_steps_per_env, + obs, + [self.env.num_actions], + ) + + return alg diff --git a/rsl_rl/runners/on_policy_runner.py b/rsl_rl/runners/on_policy_runner.py new file mode 100644 index 0000000..36f11f3 --- /dev/null +++ b/rsl_rl/runners/on_policy_runner.py @@ -0,0 +1,460 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +import statistics +import time +import torch +import warnings +from collections import deque + +import rsl_rl +from rsl_rl.algorithms import PPO +from rsl_rl.env import VecEnv +from rsl_rl.modules import ActorCritic, ActorCriticRecurrent, resolve_rnd_config, resolve_symmetry_config +from rsl_rl.utils import resolve_obs_groups, store_code_state + + +class OnPolicyRunner: + """On-policy runner for training and evaluation of actor-critic methods.""" + + def __init__(self, env: VecEnv, train_cfg: dict, log_dir: str | None = None, device="cpu"): + self.cfg = train_cfg + self.alg_cfg = train_cfg["algorithm"] + self.policy_cfg = train_cfg["policy"] + self.device = device + self.env = env + + # check if multi-gpu is enabled + self._configure_multi_gpu() + + # store training configuration + self.num_steps_per_env = self.cfg["num_steps_per_env"] + self.save_interval = self.cfg["save_interval"] + + # query observations from environment for algorithm construction + obs = self.env.get_observations() + default_sets = ["critic"] + if "rnd_cfg" in self.alg_cfg and self.alg_cfg["rnd_cfg"] is not None: + default_sets.append("rnd_state") + self.cfg["obs_groups"] = resolve_obs_groups(obs, self.cfg["obs_groups"], default_sets) + + # create the algorithm + self.alg = self._construct_algorithm(obs) + + # Decide whether to disable logging + # We only log from the process with rank 0 (main process) + self.disable_logs = self.is_distributed and self.gpu_global_rank != 0 + + # Logging + self.log_dir = log_dir + self.writer = None + self.tot_timesteps = 0 + self.tot_time = 0 + self.current_learning_iteration = 0 + self.git_status_repos = [rsl_rl.__file__] + + def learn(self, num_learning_iterations: int, init_at_random_ep_len: bool = False): # noqa: C901 + # initialize writer + self._prepare_logging_writer() + + # randomize initial episode lengths (for exploration) + if init_at_random_ep_len: + self.env.episode_length_buf = torch.randint_like( + self.env.episode_length_buf, high=int(self.env.max_episode_length) + ) + + # start learning + obs = self.env.get_observations().to(self.device) + self.train_mode() # switch to train mode (for dropout for example) + + # Book keeping + ep_infos = [] + rewbuffer = deque(maxlen=100) + lenbuffer = deque(maxlen=100) + cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + # create buffers for logging extrinsic and intrinsic rewards + if self.alg.rnd: + erewbuffer = deque(maxlen=100) + irewbuffer = deque(maxlen=100) + cur_ereward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + cur_ireward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device) + + # Ensure all parameters are in-synced + if self.is_distributed: + print(f"Synchronizing parameters for rank {self.gpu_global_rank}...") + self.alg.broadcast_parameters() + + # Start training + start_iter = self.current_learning_iteration + tot_iter = start_iter + num_learning_iterations + for it in range(start_iter, tot_iter): + start = time.time() + # Rollout + with torch.inference_mode(): + for _ in range(self.num_steps_per_env): + # Sample actions + actions = self.alg.act(obs) + # Step the environment + obs, rewards, dones, extras = self.env.step(actions.to(self.env.device)) + # Move to device + obs, rewards, dones = (obs.to(self.device), rewards.to(self.device), dones.to(self.device)) + # process the step + self.alg.process_env_step(obs, rewards, dones, extras) + # Extract intrinsic rewards (only for logging) + intrinsic_rewards = self.alg.intrinsic_rewards if self.alg.rnd else None + # book keeping + if self.log_dir is not None: + if "episode" in extras: + ep_infos.append(extras["episode"]) + elif "log" in extras: + ep_infos.append(extras["log"]) + # Update rewards + if self.alg.rnd: + cur_ereward_sum += rewards + cur_ireward_sum += intrinsic_rewards # type: ignore + cur_reward_sum += rewards + intrinsic_rewards + else: + cur_reward_sum += rewards + # Update episode length + cur_episode_length += 1 + # Clear data for completed episodes + # -- common + new_ids = (dones > 0).nonzero(as_tuple=False) + rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist()) + lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist()) + cur_reward_sum[new_ids] = 0 + cur_episode_length[new_ids] = 0 + # -- intrinsic and extrinsic rewards + if self.alg.rnd: + erewbuffer.extend(cur_ereward_sum[new_ids][:, 0].cpu().numpy().tolist()) + irewbuffer.extend(cur_ireward_sum[new_ids][:, 0].cpu().numpy().tolist()) + cur_ereward_sum[new_ids] = 0 + cur_ireward_sum[new_ids] = 0 + + stop = time.time() + collection_time = stop - start + start = stop + + # compute returns + self.alg.compute_returns(obs) + + # update policy + loss_dict = self.alg.update() + + stop = time.time() + learn_time = stop - start + self.current_learning_iteration = it + # log info + if self.log_dir is not None and not self.disable_logs: + # Log information + self.log(locals()) + # Save model + if it % self.save_interval == 0: + self.save(os.path.join(self.log_dir, f"model_{it}.pt")) + + # Clear episode infos + ep_infos.clear() + # Save code state + if it == start_iter and not self.disable_logs: + # obtain all the diff files + git_file_paths = store_code_state(self.log_dir, self.git_status_repos) + # if possible store them to wandb + if self.logger_type in ["wandb", "neptune"] and git_file_paths: + for path in git_file_paths: + self.writer.save_file(path) + + # Save the final model after training + if self.log_dir is not None and not self.disable_logs: + self.save(os.path.join(self.log_dir, f"model_{self.current_learning_iteration}.pt")) + + def log(self, locs: dict, width: int = 80, pad: int = 35): + # Compute the collection size + collection_size = self.num_steps_per_env * self.env.num_envs * self.gpu_world_size + # Update total time-steps and time + self.tot_timesteps += collection_size + self.tot_time += locs["collection_time"] + locs["learn_time"] + iteration_time = locs["collection_time"] + locs["learn_time"] + + # -- Episode info + ep_string = "" + if locs["ep_infos"]: + for key in locs["ep_infos"][0]: + infotensor = torch.tensor([], device=self.device) + for ep_info in locs["ep_infos"]: + # handle scalar and zero dimensional tensor infos + if key not in ep_info: + continue + if not isinstance(ep_info[key], torch.Tensor): + ep_info[key] = torch.Tensor([ep_info[key]]) + if len(ep_info[key].shape) == 0: + ep_info[key] = ep_info[key].unsqueeze(0) + infotensor = torch.cat((infotensor, ep_info[key].to(self.device))) + value = torch.mean(infotensor) + # log to logger and terminal + if "/" in key: + self.writer.add_scalar(key, value, locs["it"]) + ep_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + else: + self.writer.add_scalar("Episode/" + key, value, locs["it"]) + ep_string += f"""{f'Mean episode {key}:':>{pad}} {value:.4f}\n""" + + mean_std = self.alg.policy.action_std.mean() + fps = int(collection_size / (locs["collection_time"] + locs["learn_time"])) + + # -- Losses + for key, value in locs["loss_dict"].items(): + self.writer.add_scalar(f"Loss/{key}", value, locs["it"]) + self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"]) + + # -- Policy + self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"]) + + # -- Performance + self.writer.add_scalar("Perf/total_fps", fps, locs["it"]) + self.writer.add_scalar("Perf/collection time", locs["collection_time"], locs["it"]) + self.writer.add_scalar("Perf/learning_time", locs["learn_time"], locs["it"]) + + # -- Training + if len(locs["rewbuffer"]) > 0: + # separate logging for intrinsic and extrinsic rewards + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.writer.add_scalar("Rnd/mean_extrinsic_reward", statistics.mean(locs["erewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/mean_intrinsic_reward", statistics.mean(locs["irewbuffer"]), locs["it"]) + self.writer.add_scalar("Rnd/weight", self.alg.rnd.weight, locs["it"]) + # everything else + self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"]) + self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"]) + if self.logger_type != "wandb": # wandb does not support non-integer x-axis logging + self.writer.add_scalar("Train/mean_reward/time", statistics.mean(locs["rewbuffer"]), self.tot_time) + self.writer.add_scalar( + "Train/mean_episode_length/time", statistics.mean(locs["lenbuffer"]), self.tot_time + ) + + str = f" \033[1m Learning iteration {locs['it']}/{locs['tot_iter']} \033[0m " + + if len(locs["rewbuffer"]) > 0: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + ) + # -- Losses + for key, value in locs["loss_dict"].items(): + log_string += f"""{f'Mean {key} loss:':>{pad}} {value:.4f}\n""" + # -- Rewards + if hasattr(self.alg, "rnd") and self.alg.rnd: + log_string += ( + f"""{'Mean extrinsic reward:':>{pad}} {statistics.mean(locs['erewbuffer']):.2f}\n""" + f"""{'Mean intrinsic reward:':>{pad}} {statistics.mean(locs['irewbuffer']):.2f}\n""" + ) + log_string += f"""{'Mean reward:':>{pad}} {statistics.mean(locs['rewbuffer']):.2f}\n""" + # -- episode info + log_string += f"""{'Mean episode length:':>{pad}} {statistics.mean(locs['lenbuffer']):.2f}\n""" + else: + log_string = ( + f"""{'#' * width}\n""" + f"""{str.center(width, ' ')}\n\n""" + f"""{'Computation:':>{pad}} {fps:.0f} steps/s (collection: {locs[ + 'collection_time']:.3f}s, learning {locs['learn_time']:.3f}s)\n""" + f"""{'Mean action noise std:':>{pad}} {mean_std.item():.2f}\n""" + ) + for key, value in locs["loss_dict"].items(): + log_string += f"""{f'{key}:':>{pad}} {value:.4f}\n""" + + log_string += ep_string + log_string += ( + f"""{'-' * width}\n""" + f"""{'Total timesteps:':>{pad}} {self.tot_timesteps}\n""" + f"""{'Iteration time:':>{pad}} {iteration_time:.2f}s\n""" + f"""{'Time elapsed:':>{pad}} {time.strftime("%H:%M:%S", time.gmtime(self.tot_time))}\n""" + f"""{'ETA:':>{pad}} {time.strftime( + "%H:%M:%S", + time.gmtime( + self.tot_time / (locs['it'] - locs['start_iter'] + 1) + * (locs['start_iter'] + locs['num_learning_iterations'] - locs['it']) + ) + )}\n""" + ) + print(log_string) + + def save(self, path: str, infos=None): + # -- Save model + saved_dict = { + "model_state_dict": self.alg.policy.state_dict(), + "optimizer_state_dict": self.alg.optimizer.state_dict(), + "iter": self.current_learning_iteration, + "infos": infos, + } + # -- Save RND model if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + saved_dict["rnd_state_dict"] = self.alg.rnd.state_dict() + saved_dict["rnd_optimizer_state_dict"] = self.alg.rnd_optimizer.state_dict() + torch.save(saved_dict, path) + + # upload model to external logging service + if self.logger_type in ["neptune", "wandb"] and not self.disable_logs: + self.writer.save_model(path, self.current_learning_iteration) + + def load(self, path: str, load_optimizer: bool = True, map_location: str | None = None): + loaded_dict = torch.load(path, weights_only=False, map_location=map_location) + # -- Load model + resumed_training = self.alg.policy.load_state_dict(loaded_dict["model_state_dict"]) + # -- Load RND model if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.load_state_dict(loaded_dict["rnd_state_dict"]) + # -- load optimizer if used + if load_optimizer and resumed_training: + # -- algorithm optimizer + self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"]) + # -- RND optimizer if used + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd_optimizer.load_state_dict(loaded_dict["rnd_optimizer_state_dict"]) + # -- load current learning iteration + if resumed_training: + self.current_learning_iteration = loaded_dict["iter"] + return loaded_dict["infos"] + + def get_inference_policy(self, device=None): + self.eval_mode() # switch to evaluation mode (dropout for example) + if device is not None: + self.alg.policy.to(device) + return self.alg.policy.act_inference + + def train_mode(self): + # -- PPO + self.alg.policy.train() + # -- RND + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.train() + + def eval_mode(self): + # -- PPO + self.alg.policy.eval() + # -- RND + if hasattr(self.alg, "rnd") and self.alg.rnd: + self.alg.rnd.eval() + + def add_git_repo_to_log(self, repo_file_path): + self.git_status_repos.append(repo_file_path) + + """ + Helper functions. + """ + + def _configure_multi_gpu(self): + """Configure multi-gpu training.""" + # check if distributed training is enabled + self.gpu_world_size = int(os.getenv("WORLD_SIZE", "1")) + self.is_distributed = self.gpu_world_size > 1 + + # if not distributed training, set local and global rank to 0 and return + if not self.is_distributed: + self.gpu_local_rank = 0 + self.gpu_global_rank = 0 + self.multi_gpu_cfg = None + return + + # get rank and world size + self.gpu_local_rank = int(os.getenv("LOCAL_RANK", "0")) + self.gpu_global_rank = int(os.getenv("RANK", "0")) + + # make a configuration dictionary + self.multi_gpu_cfg = { + "global_rank": self.gpu_global_rank, # rank of the main process + "local_rank": self.gpu_local_rank, # rank of the current process + "world_size": self.gpu_world_size, # total number of processes + } + + # check if user has device specified for local rank + if self.device != f"cuda:{self.gpu_local_rank}": + raise ValueError( + f"Device '{self.device}' does not match expected device for local rank '{self.gpu_local_rank}'." + ) + # validate multi-gpu configuration + if self.gpu_local_rank >= self.gpu_world_size: + raise ValueError( + f"Local rank '{self.gpu_local_rank}' is greater than or equal to world size '{self.gpu_world_size}'." + ) + if self.gpu_global_rank >= self.gpu_world_size: + raise ValueError( + f"Global rank '{self.gpu_global_rank}' is greater than or equal to world size '{self.gpu_world_size}'." + ) + + # initialize torch distributed + torch.distributed.init_process_group(backend="nccl", rank=self.gpu_global_rank, world_size=self.gpu_world_size) + # set device to the local rank + torch.cuda.set_device(self.gpu_local_rank) + + def _construct_algorithm(self, obs) -> PPO: + """Construct the actor-critic algorithm.""" + # resolve RND config + self.alg_cfg = resolve_rnd_config(self.alg_cfg, obs, self.cfg["obs_groups"], self.env) + + # resolve symmetry config + self.alg_cfg = resolve_symmetry_config(self.alg_cfg, self.env) + + # resolve deprecated normalization config + if self.cfg.get("empirical_normalization") is not None: + warnings.warn( + "The `empirical_normalization` parameter is deprecated. Please set `actor_obs_normalization` and " + "`critic_obs_normalization` as part of the `policy` configuration instead.", + DeprecationWarning, + ) + if self.policy_cfg.get("actor_obs_normalization") is None: + self.policy_cfg["actor_obs_normalization"] = self.cfg["empirical_normalization"] + if self.policy_cfg.get("critic_obs_normalization") is None: + self.policy_cfg["critic_obs_normalization"] = self.cfg["empirical_normalization"] + + # initialize the actor-critic + actor_critic_class = eval(self.policy_cfg.pop("class_name")) + actor_critic: ActorCritic | ActorCriticRecurrent = actor_critic_class( + obs, self.cfg["obs_groups"], self.env.num_actions, **self.policy_cfg + ).to(self.device) + + # initialize the algorithm + alg_class = eval(self.alg_cfg.pop("class_name")) + alg: PPO = alg_class(actor_critic, device=self.device, **self.alg_cfg, multi_gpu_cfg=self.multi_gpu_cfg) + + # initialize the storage + alg.init_storage( + "rl", + self.env.num_envs, + self.num_steps_per_env, + obs, + [self.env.num_actions], + ) + + return alg + + def _prepare_logging_writer(self): + """Prepares the logging writers.""" + if self.log_dir is not None and self.writer is None and not self.disable_logs: + # Launch either Tensorboard or Neptune & Tensorboard summary writer(s), default: Tensorboard. + self.logger_type = self.cfg.get("logger", "tensorboard") + self.logger_type = self.logger_type.lower() + + if self.logger_type == "neptune": + from rsl_rl.utils.neptune_utils import NeptuneSummaryWriter + + self.writer = NeptuneSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "wandb": + from rsl_rl.utils.wandb_utils import WandbSummaryWriter + + self.writer = WandbSummaryWriter(log_dir=self.log_dir, flush_secs=10, cfg=self.cfg) + self.writer.log_config(self.env.cfg, self.cfg, self.alg_cfg, self.policy_cfg) + elif self.logger_type == "tensorboard": + from torch.utils.tensorboard import SummaryWriter + + self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10) + else: + raise ValueError("Logger type not found. Please choose 'neptune', 'wandb' or 'tensorboard'.") diff --git a/rsl_rl/storage/__init__.py b/rsl_rl/storage/__init__.py new file mode 100644 index 0000000..1624330 --- /dev/null +++ b/rsl_rl/storage/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Implementation of transitions storage for RL-agent.""" + +from .rollout_storage import RolloutStorage +from .replay_buffer_multi import ReplayBufferMulti +__all__ = ["RolloutStorage", "ReplayBufferMulti"] diff --git a/rsl_rl/storage/__pycache__/__init__.cpython-312.pyc b/rsl_rl/storage/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..12b3a67 Binary files /dev/null and b/rsl_rl/storage/__pycache__/__init__.cpython-312.pyc differ diff --git a/rsl_rl/storage/__pycache__/replay_buffer_multi.cpython-312.pyc b/rsl_rl/storage/__pycache__/replay_buffer_multi.cpython-312.pyc new file mode 100644 index 0000000..8b19914 Binary files /dev/null and b/rsl_rl/storage/__pycache__/replay_buffer_multi.cpython-312.pyc differ diff --git a/rsl_rl/storage/__pycache__/rollout_storage.cpython-312.pyc b/rsl_rl/storage/__pycache__/rollout_storage.cpython-312.pyc new file mode 100644 index 0000000..b83a0f6 Binary files /dev/null and b/rsl_rl/storage/__pycache__/rollout_storage.cpython-312.pyc differ diff --git a/rsl_rl/storage/replay_buffer_multi.py b/rsl_rl/storage/replay_buffer_multi.py new file mode 100644 index 0000000..6462b8e --- /dev/null +++ b/rsl_rl/storage/replay_buffer_multi.py @@ -0,0 +1,38 @@ +import torch +import numpy as np + + +class ReplayBufferMulti: + """Fixed-size buffer to store experience tuples.""" + + def __init__(self, obs_dim, buffer_size, num_amp_frames, device): + """Initialize a ReplayBuffer object. + Arguments: + buffer_size (int): maximum size of buffer + """ + self.states = torch.zeros(buffer_size, num_amp_frames, obs_dim).to(device) + self.num_amp_frames = num_amp_frames + self.buffer_size = buffer_size + self.device = device + + self.step = 0 + self.num_samples = 0 + + def insert(self, states): + """Add new states to memory.""" + num_states = states.shape[0] + start_idx = self.step + end_idx = self.step + num_states + if end_idx > self.buffer_size: + self.states[self.step:self.buffer_size] = states[:self.buffer_size - self.step] + self.states[:end_idx - self.buffer_size] = states[self.buffer_size - self.step:] + else: + self.states[start_idx:end_idx] = states + + self.num_samples = min(self.buffer_size, max(end_idx, self.num_samples)) + self.step = (self.step + num_states) % self.buffer_size + + def feed_forward_generator(self, num_mini_batch, mini_batch_size): + for _ in range(num_mini_batch): + sample_idxs = np.random.choice(self.num_samples, size=mini_batch_size) + yield (self.states[sample_idxs].to(self.device)) diff --git a/rsl_rl/storage/rollout_storage.py b/rsl_rl/storage/rollout_storage.py new file mode 100644 index 0000000..e9309b3 --- /dev/null +++ b/rsl_rl/storage/rollout_storage.py @@ -0,0 +1,260 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import torch +from tensordict import TensorDict + +from rsl_rl.utils import split_and_pad_trajectories + + +class RolloutStorage: + class Transition: + def __init__(self): + self.observations = None + self.actions = None + self.privileged_actions = None + self.rewards = None + self.dones = None + self.values = None + self.actions_log_prob = None + self.action_mean = None + self.action_sigma = None + self.hidden_states = None + + def clear(self): + self.__init__() + + def __init__( + self, + training_type, + num_envs, + num_transitions_per_env, + obs, + actions_shape, + device="cpu", + ): + # store inputs + self.training_type = training_type + self.device = device + self.num_transitions_per_env = num_transitions_per_env + self.num_envs = num_envs + self.actions_shape = actions_shape + + # Core + self.observations = TensorDict( + {key: torch.zeros(num_transitions_per_env, *value.shape, device=device) for key, value in obs.items()}, + batch_size=[num_transitions_per_env, num_envs], + device=self.device, + ) + self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte() + + # for distillation + if training_type == "distillation": + self.privileged_actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + + # for reinforcement learning + if training_type == "rl": + self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) + self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) + + # For RNN networks + self.saved_hidden_states_a = None + self.saved_hidden_states_c = None + + # counter for the number of transitions stored + self.step = 0 + + def add_transitions(self, transition: Transition): + # check if the transition is valid + if self.step >= self.num_transitions_per_env: + raise OverflowError("Rollout buffer overflow! You should call clear() before adding new transitions.") + + # Core + self.observations[self.step].copy_(transition.observations) + self.actions[self.step].copy_(transition.actions) + self.rewards[self.step].copy_(transition.rewards.view(-1, 1)) + self.dones[self.step].copy_(transition.dones.view(-1, 1)) + + # for distillation + if self.training_type == "distillation": + self.privileged_actions[self.step].copy_(transition.privileged_actions) + + # for reinforcement learning + if self.training_type == "rl": + self.values[self.step].copy_(transition.values) + self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1)) + self.mu[self.step].copy_(transition.action_mean) + self.sigma[self.step].copy_(transition.action_sigma) + + # For RNN networks + self._save_hidden_states(transition.hidden_states) + + # increment the counter + self.step += 1 + + def _save_hidden_states(self, hidden_states): + if hidden_states is None or hidden_states == (None, None): + return + # make a tuple out of GRU hidden state sto match the LSTM format + hid_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],) + hid_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],) + # initialize if needed + if self.saved_hidden_states_a is None: + self.saved_hidden_states_a = [ + torch.zeros(self.observations.shape[0], *hid_a[i].shape, device=self.device) for i in range(len(hid_a)) + ] + self.saved_hidden_states_c = [ + torch.zeros(self.observations.shape[0], *hid_c[i].shape, device=self.device) for i in range(len(hid_c)) + ] + # copy the states + for i in range(len(hid_a)): + self.saved_hidden_states_a[i][self.step].copy_(hid_a[i]) + self.saved_hidden_states_c[i][self.step].copy_(hid_c[i]) + + def clear(self): + self.step = 0 + + def compute_returns(self, last_values, gamma, lam, normalize_advantage: bool = True): + advantage = 0 + for step in reversed(range(self.num_transitions_per_env)): + # if we are at the last step, bootstrap the return value + if step == self.num_transitions_per_env - 1: + next_values = last_values + else: + next_values = self.values[step + 1] + # 1 if we are not in a terminal state, 0 otherwise + next_is_not_terminal = 1.0 - self.dones[step].float() + # TD error: r_t + gamma * V(s_{t+1}) - V(s_t) + delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step] + # Advantage: A(s_t, a_t) = delta_t + gamma * lambda * A(s_{t+1}, a_{t+1}) + advantage = delta + next_is_not_terminal * gamma * lam * advantage + # Return: R_t = A(s_t, a_t) + V(s_t) + self.returns[step] = advantage + self.values[step] + + # Compute the advantages + self.advantages = self.returns - self.values + # Normalize the advantages if flag is set + # This is to prevent double normalization (i.e. if per minibatch normalization is used) + if normalize_advantage: + self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8) + + # for distillation + def generator(self): + if self.training_type != "distillation": + raise ValueError("This function is only available for distillation training.") + + for i in range(self.num_transitions_per_env): + yield self.observations[i], self.actions[i], self.privileged_actions[i], self.dones[i] + + # for reinforcement learning with feedforward networks + def mini_batch_generator(self, num_mini_batches, num_epochs=8): + if self.training_type != "rl": + raise ValueError("This function is only available for reinforcement learning training.") + batch_size = self.num_envs * self.num_transitions_per_env + mini_batch_size = batch_size // num_mini_batches + indices = torch.randperm(num_mini_batches * mini_batch_size, requires_grad=False, device=self.device) + + # Core + observations = self.observations.flatten(0, 1) + actions = self.actions.flatten(0, 1) + values = self.values.flatten(0, 1) + returns = self.returns.flatten(0, 1) + + # For PPO + old_actions_log_prob = self.actions_log_prob.flatten(0, 1) + advantages = self.advantages.flatten(0, 1) + old_mu = self.mu.flatten(0, 1) + old_sigma = self.sigma.flatten(0, 1) + + for epoch in range(num_epochs): + for i in range(num_mini_batches): + # Select the indices for the mini-batch + start = i * mini_batch_size + end = (i + 1) * mini_batch_size + batch_idx = indices[start:end] + + # Create the mini-batch + # -- Core + obs_batch = observations[batch_idx] + actions_batch = actions[batch_idx] + + # -- For PPO + target_values_batch = values[batch_idx] + returns_batch = returns[batch_idx] + old_actions_log_prob_batch = old_actions_log_prob[batch_idx] + advantages_batch = advantages[batch_idx] + old_mu_batch = old_mu[batch_idx] + old_sigma_batch = old_sigma[batch_idx] + + # yield the mini-batch + yield obs_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( + None, + None, + ), None + + # for reinfrocement learning with recurrent networks + def recurrent_mini_batch_generator(self, num_mini_batches, num_epochs=8): + if self.training_type != "rl": + raise ValueError("This function is only available for reinforcement learning training.") + padded_obs_trajectories, trajectory_masks = split_and_pad_trajectories(self.observations, self.dones) + + mini_batch_size = self.num_envs // num_mini_batches + for ep in range(num_epochs): + first_traj = 0 + for i in range(num_mini_batches): + start = i * mini_batch_size + stop = (i + 1) * mini_batch_size + + dones = self.dones.squeeze(-1) + last_was_done = torch.zeros_like(dones, dtype=torch.bool) + last_was_done[1:] = dones[:-1] + last_was_done[0] = True + trajectories_batch_size = torch.sum(last_was_done[:, start:stop]) + last_traj = first_traj + trajectories_batch_size + + masks_batch = trajectory_masks[:, first_traj:last_traj] + obs_batch = padded_obs_trajectories[:, first_traj:last_traj] + actions_batch = self.actions[:, start:stop] + old_mu_batch = self.mu[:, start:stop] + old_sigma_batch = self.sigma[:, start:stop] + returns_batch = self.returns[:, start:stop] + advantages_batch = self.advantages[:, start:stop] + values_batch = self.values[:, start:stop] + old_actions_log_prob_batch = self.actions_log_prob[:, start:stop] + + # reshape to [num_envs, time, num layers, hidden dim] (original shape: [time, num_layers, num_envs, hidden_dim]) + # then take only time steps after dones (flattens num envs and time dimensions), + # take a batch of trajectories and finally reshape back to [num_layers, batch, hidden_dim] + last_was_done = last_was_done.permute(1, 0) + hid_a_batch = [ + saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] + .transpose(1, 0) + .contiguous() + for saved_hidden_states in self.saved_hidden_states_a + ] + hid_c_batch = [ + saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] + .transpose(1, 0) + .contiguous() + for saved_hidden_states in self.saved_hidden_states_c + ] + # remove the tuple for GRU + hid_a_batch = hid_a_batch[0] if len(hid_a_batch) == 1 else hid_a_batch + hid_c_batch = hid_c_batch[0] if len(hid_c_batch) == 1 else hid_c_batch + + yield obs_batch, actions_batch, values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( + hid_a_batch, + hid_c_batch, + ), masks_batch + + first_traj = last_traj diff --git a/rsl_rl/utils/__init__.py b/rsl_rl/utils/__init__.py new file mode 100644 index 0000000..f5781f1 --- /dev/null +++ b/rsl_rl/utils/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Helper functions.""" + +from .utils import * +from .motion_loader_g1 import G1_AMPLoader + +__all__ = [ + "G1_AMPLoader", +] \ No newline at end of file diff --git a/rsl_rl/utils/__pycache__/__init__.cpython-312.pyc b/rsl_rl/utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..b582fb3 Binary files /dev/null and b/rsl_rl/utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/rsl_rl/utils/__pycache__/motion_loader_g1.cpython-312.pyc b/rsl_rl/utils/__pycache__/motion_loader_g1.cpython-312.pyc new file mode 100644 index 0000000..499d556 Binary files /dev/null and b/rsl_rl/utils/__pycache__/motion_loader_g1.cpython-312.pyc differ diff --git a/rsl_rl/utils/__pycache__/motion_util.cpython-312.pyc b/rsl_rl/utils/__pycache__/motion_util.cpython-312.pyc new file mode 100644 index 0000000..858b004 Binary files /dev/null and b/rsl_rl/utils/__pycache__/motion_util.cpython-312.pyc differ diff --git a/rsl_rl/utils/__pycache__/pose3d.cpython-312.pyc b/rsl_rl/utils/__pycache__/pose3d.cpython-312.pyc new file mode 100644 index 0000000..39d0853 Binary files /dev/null and b/rsl_rl/utils/__pycache__/pose3d.cpython-312.pyc differ diff --git a/rsl_rl/utils/__pycache__/utils.cpython-312.pyc b/rsl_rl/utils/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000..1c2e1a3 Binary files /dev/null and b/rsl_rl/utils/__pycache__/utils.cpython-312.pyc differ diff --git a/rsl_rl/utils/motion_loader_g1.py b/rsl_rl/utils/motion_loader_g1.py new file mode 100644 index 0000000..677c630 --- /dev/null +++ b/rsl_rl/utils/motion_loader_g1.py @@ -0,0 +1,388 @@ +import os +from os.path import join as pjoin +import glob +import json +import logging + +import torch +import numpy as np +from pybullet_utils import transformations + +from rsl_rl.utils import motion_util + +_EPS = np.finfo(float).eps * 4.0 +def quaternion_slerp(q0, q1, fraction, spin=0, shortestpath=True): + """Batch quaternion spherical linear interpolation.""" + + out = torch.zeros_like(q0) + + zero_mask = torch.isclose(fraction, torch.zeros_like(fraction)).squeeze() + ones_mask = torch.isclose(fraction, torch.ones_like(fraction)).squeeze() + out[zero_mask] = q0[zero_mask] + out[ones_mask] = q1[ones_mask] + + d = torch.sum(q0 * q1, dim=-1, keepdim=True) + dist_mask = (torch.abs(torch.abs(d) - 1.0) < _EPS).squeeze() + out[dist_mask] = q0[dist_mask] + + if shortestpath: + d_old = torch.clone(d) + d = torch.where(d_old < 0, -d, d) + q1 = torch.where(d_old < 0, -q1, q1) + + angle = torch.acos(d) + spin * torch.pi + angle_mask = (torch.abs(angle) < _EPS).squeeze() + out[angle_mask] = q0[angle_mask] + + final_mask = torch.logical_or(zero_mask, ones_mask) + final_mask = torch.logical_or(final_mask, dist_mask) + final_mask = torch.logical_or(final_mask, angle_mask) + final_mask = torch.logical_not(final_mask) + + isin = 1.0 / angle + q0 *= torch.sin((1.0 - fraction) * angle) * isin + q1 *= torch.sin(fraction * angle) * isin + q0 += q1 + out[final_mask] = q0[final_mask] + return out + + +class G1_AMPLoader: + + def __init__( + self, + device, + time_between_frames, + motion_files, + preload_transitions=False, + num_preload_transitions=1000000, + num_frames=5, + ): + """Expert dataset provides AMP observations from Dog mocap dataset. + + time_between_frames: Amount of time in seconds between transition. + """ + self.device = device + self.time_between_frames = time_between_frames + self.num_frames = num_frames + + # Values to store for each trajectory. + self.trajectories = [] + self.trajectories_full = [] + self.trajectory_names = [] + self.trajectory_idxs = [] + self.trajectory_lens = [] # Traj length in seconds. + self.trajectory_weights = [] + self.trajectory_frame_durations = [] + self.trajectory_num_frames = [] + self.motion_dir = motion_files + # import ipdb; ipdb.set_trace() + for i, motion_file in enumerate(os.listdir(motion_files)): + self.trajectory_names.append(motion_file) + motion_path = pjoin(motion_files, motion_file) + motion_data = np.load(motion_path, allow_pickle=True) + motion_data_processed = np.zeros((motion_data.shape[0],36)) + + for f_i in range(motion_data.shape[0]): + motion_data_processed[f_i, :3] = motion_data[f_i, :3] # base pos + motion_data_processed[f_i, 3:7] = motion_data[f_i, 3:7] # base quat (wxyz) + motion_data_processed[f_i, 7:35] = motion_data[f_i, 7:35] # base vel + ''' + NOTE The order of motion_data_processed is + base pos 0:3, + base quat 3:7, wxyz + dof pos 7:36, (mujoco joint order) + ''' + self.trajectories.append(torch.tensor( + motion_data_processed[:, 7:], + dtype=torch.float32, + device=self.device + )) + + self.trajectories_full.append(torch.tensor( + motion_data_processed, + dtype=torch.float32, + device=self.device + )) + + self.trajectory_idxs.append(i) + self.trajectory_weights.append(1 / len(os.listdir(motion_files))) + frame_duration = 1 / 50 + + self.trajectory_frame_durations.append(frame_duration) + traj_len = (motion_data_processed.shape[0] - 1) * frame_duration # seconds + self.trajectory_lens.append(traj_len) + self.trajectory_num_frames.append(float(motion_data_processed.shape[0])) + print(f"Loaded {traj_len}s. motion from {motion_file}.") + + # Trajectory weights are used to sample some trajectories more than others. + self.trajectory_weights = np.array(self.trajectory_weights) / np.sum(self.trajectory_weights) + self.trajectory_frame_durations = np.array(self.trajectory_frame_durations) + self.trajectory_lens = np.array(self.trajectory_lens) + self.trajectory_num_frames = np.array(self.trajectory_num_frames) + + # Preload transitions. + self.preload_transitions = preload_transitions + if self.preload_transitions: + print(f'Preloading {num_preload_transitions} transitions') + + traj_idxs = self.weighted_traj_idx_sample_batch(num_preload_transitions) + times = self.traj_time_sample_batch(traj_idxs) + self.preloaded_s_prior = self.get_full_frame_at_time_batch(traj_idxs, times - self.time_between_frames) + self.preloaded_s = self.get_full_frame_at_time_batch(traj_idxs, times) + self.preloaded_s_next = self.get_full_frame_at_time_batch(traj_idxs, times + self.time_between_frames) + print(f'Finished preloading') + + # 预加载多帧数据 + self.preloaded_frames = [] + for i in range(self.num_frames): + frame_time = times + (i - (self.num_frames - 2)) * self.time_between_frames + full_frame = self.get_full_frame_at_time_batch(traj_idxs, frame_time) + # 预处理：提前提取并连接需要的列（7:26 和 29:33），避免每次生成时重复切片 + processed_frame = torch.cat([ + full_frame[:, 7:26], + full_frame[:, 29:33] + ], dim=-1) + self.preloaded_frames.append(processed_frame) + print(f'Finished preloading multiple frames') + + self.all_trajectories_full = torch.vstack(self.trajectories_full) + + def weighted_traj_idx_sample(self): + """Get traj idx via weighted sampling.""" + return np.random.choice( + self.trajectory_idxs, p=self.trajectory_weights) + + def weighted_traj_idx_sample_batch(self, size): + """Batch sample traj idxs.""" + return np.random.choice( + self.trajectory_idxs, size=size, p=self.trajectory_weights, + replace=True) + + def traj_time_sample(self, traj_idx): + """Sample random time for traj.""" + subst = self.time_between_frames + self.trajectory_frame_durations[traj_idx] + return max( + 0, (self.trajectory_lens[traj_idx] * np.random.uniform() - subst)) + + def traj_time_sample_batch(self, traj_idxs): + """Sample random time for multiple trajectories.""" + subst = self.time_between_frames + self.trajectory_frame_durations[traj_idxs] + time_samples = self.trajectory_lens[traj_idxs] * np.random.uniform(size=len(traj_idxs)) - subst + return np.maximum(np.zeros_like(time_samples), time_samples) + + def slerp(self, val0, val1, blend): + return (1.0 - blend) * val0 + blend * val1 + + def get_trajectory(self, traj_idx): + """Returns trajectory of AMP observations.""" + return self.trajectories_full[traj_idx] + + def get_frame_at_time(self, traj_idx, time): + """Returns frame for the given trajectory at the specified time.""" + p = float(time) / self.trajectory_lens[traj_idx] + n = self.trajectories[traj_idx].shape[0] + idx_low, idx_high = int(np.floor(p * n)), int(np.ceil(p * n)) + frame_start = self.trajectories[traj_idx][idx_low] + frame_end = self.trajectories[traj_idx][idx_high] + blend = p * n - idx_low + return self.slerp(frame_start, frame_end, blend) + + def get_frame_at_time_batch(self, traj_idxs, times): + """Returns frame for the given trajectory at the specified time.""" + p = times / self.trajectory_lens[traj_idxs] + n = self.trajectory_num_frames[traj_idxs] + idx_low, idx_high = np.floor(p * n).astype(np.int32), np.ceil(p * n).astype(np.int32) + all_frame_starts = torch.zeros(len(traj_idxs), self.observation_dim, device=self.device) + all_frame_ends = torch.zeros(len(traj_idxs), self.observation_dim, device=self.device) + for traj_idx in set(traj_idxs): + trajectory = self.trajectories[traj_idx] + traj_mask = traj_idxs == traj_idx + all_frame_starts[traj_mask] = trajectory[idx_low[traj_mask]] + all_frame_ends[traj_mask] = trajectory[idx_high[traj_mask]] + blend = torch.tensor(p * n - idx_low, device=self.device, dtype=torch.float32).unsqueeze(-1) + return self.slerp(all_frame_starts, all_frame_ends, blend) + + def get_full_frame_at_time(self, traj_idx, time): + """Returns full frame for the given trajectory at the specified time.""" + p = float(time) / self.trajectory_lens[traj_idx] + n = self.trajectories_full[traj_idx].shape[0] + idx_low, idx_high = int(np.floor(p * n)), int(np.ceil(p * n)) + frame_start = self.trajectories_full[traj_idx][idx_low] + frame_end = self.trajectories_full[traj_idx][idx_high] + blend = p * n - idx_low + print(idx_low, idx_high) + return self.blend_frame_pose(frame_start, frame_end, blend) + + def get_full_frame_at_time_batch(self, traj_idxs, times): + p = times / self.trajectory_lens[traj_idxs] + n = self.trajectory_num_frames[traj_idxs] + idx_low, idx_high = np.floor(p * n).astype(np.int32), np.ceil(p * n).astype(np.int32) + all_frame_pos_starts = torch.zeros(len(traj_idxs), 3, device=self.device) + all_frame_pos_ends = torch.zeros(len(traj_idxs), 3, device=self.device) + all_frame_rot_starts = torch.zeros(len(traj_idxs), 4, device=self.device) + all_frame_rot_ends = torch.zeros(len(traj_idxs), 4, device=self.device) + all_frame_amp_starts = torch.zeros(len(traj_idxs), 29, device=self.device) + all_frame_amp_ends = torch.zeros(len(traj_idxs), 29, device=self.device) + for traj_idx in set(traj_idxs): + trajectory = self.trajectories_full[traj_idx] + traj_mask = traj_idxs == traj_idx + all_frame_pos_starts[traj_mask] = G1_AMPLoader.get_root_pos_batch(trajectory[idx_low[traj_mask]]) + all_frame_pos_ends[traj_mask] = G1_AMPLoader.get_root_pos_batch(trajectory[idx_high[traj_mask]]) + all_frame_rot_starts[traj_mask] = G1_AMPLoader.get_root_rot_batch(trajectory[idx_low[traj_mask]]) + all_frame_rot_ends[traj_mask] = G1_AMPLoader.get_root_rot_batch(trajectory[idx_high[traj_mask]]) + all_frame_amp_starts[traj_mask] = trajectory[idx_low[traj_mask]][:, 7:36] # base vel3+ang3, dof vel23+ang23 + all_frame_amp_ends[traj_mask] = trajectory[idx_high[traj_mask]][:, 7:36] # base vel3+ang3, dof vel23+ang23 + blend = torch.tensor(p * n - idx_low, device=self.device, dtype=torch.float32).unsqueeze(-1) + pos_blend = self.slerp(all_frame_pos_starts, all_frame_pos_ends, blend) + rot_blend = quaternion_slerp(all_frame_rot_starts, all_frame_rot_ends, blend) + amp_blend = self.slerp(all_frame_amp_starts, all_frame_amp_ends, blend) + return torch.cat([pos_blend, rot_blend, amp_blend], dim=-1) + + def get_frame(self): + """Returns random frame.""" + traj_idx = self.weighted_traj_idx_sample() + sampled_time = self.traj_time_sample(traj_idx) + return self.get_frame_at_time(traj_idx, sampled_time) + + def get_full_frame(self): + """Returns random full frame.""" + traj_idx = self.weighted_traj_idx_sample() + sampled_time = self.traj_time_sample(traj_idx) + return self.get_full_frame_at_time(traj_idx, sampled_time) + + def get_full_frame_batch(self, num_frames): + if self.preload_transitions: + idxs = np.random.choice( + self.preloaded_s.shape[0], size=num_frames) + return self.preloaded_s[idxs] + else: + traj_idxs = self.weighted_traj_idx_sample_batch(num_frames) + times = self.traj_time_sample_batch(traj_idxs) + return self.get_full_frame_at_time_batch(traj_idxs, times) + + def blend_frame_pose(self, frame0, frame1, blend): + """Linearly interpolate between two frames, including orientation. + + Args: + frame0: First frame to be blended corresponds to (blend = 0). + frame1: Second frame to be blended corresponds to (blend = 1). + blend: Float between [0, 1], specifying the interpolation between + the two frames. + Returns: + An interpolation of the two frames. + """ + root_pos0, root_pos1 = G1_AMPLoader.get_root_pos(frame0), G1_AMPLoader.get_root_pos(frame1) + root_rot0, root_rot1 = G1_AMPLoader.get_root_rot(frame0), G1_AMPLoader.get_root_rot(frame1) + joints0, joints1 = G1_AMPLoader.get_joint_pose(frame0), G1_AMPLoader.get_joint_pose(frame1) + # tar_toe_pos_0, tar_toe_pos_1 = G1_AMPLoader.get_tar_toe_pos_local(frame0), G1_AMPLoader.get_tar_toe_pos_local(frame1) + linear_vel_0, linear_vel_1 = G1_AMPLoader.get_linear_vel(frame0), G1_AMPLoader.get_linear_vel(frame1) + angular_vel_0, angular_vel_1 = G1_AMPLoader.get_angular_vel(frame0), G1_AMPLoader.get_angular_vel(frame1) + joint_vel_0, joint_vel_1 = G1_AMPLoader.get_joint_vel(frame0), G1_AMPLoader.get_joint_vel(frame1) + + blend_root_pos = self.slerp(root_pos0, root_pos1, blend) + blend_root_rot = transformations.quaternion_slerp(root_rot0.cpu().numpy(), root_rot1.cpu().numpy(), blend) + blend_root_rot = torch.tensor(motion_util.standardize_quaternion(blend_root_rot),dtype=torch.float32, device=self.device) + blend_joints = self.slerp(joints0, joints1, blend) + # blend_tar_toe_pos = self.slerp(tar_toe_pos_0, tar_toe_pos_1, blend) + blend_linear_vel = self.slerp(linear_vel_0, linear_vel_1, blend) + blend_angular_vel = self.slerp(angular_vel_0, angular_vel_1, blend) + blend_joints_vel = self.slerp(joint_vel_0, joint_vel_1, blend) + + # return + # torch.cat([ + # blend_root_pos, blend_root_rot, blend_linear_vel, blend_angular_vel, blend_joints, blend_joints_vel]) + return torch.cat([blend_root_pos, blend_root_rot, blend_linear_vel, blend_angular_vel, blend_joints]) + + def feed_forward_generator_23dof_multi(self, num_mini_batch, mini_batch_size): + """Generates a batch of AMP transitions.""" + # import ipdb; ipdb.set_trace() + for _ in range(num_mini_batch): + if self.preload_transitions: + idxs = np.random.choice(self.preloaded_s.shape[0], size=mini_batch_size) + + frames = [] + for i in range(self.num_frames): + # 数据已在预加载时预处理，直接索引即可 + s = self.preloaded_frames[i][idxs] + frames.append(s) + else: + NotImplementedError('preload transition') + yield torch.stack(frames, dim=1) # [batch, num_frames, 16] + + + + + def quaternion_to_euler_array(self, quat): + # Ensure quaternion is in the correct format [x, y, z, w] + x, y, z, w =quat + + # Roll (x-axis rotation) + t0 = +2.0 * (w * x + y * z) + t1 = +1.0 - 2.0 * (x * x + y * y) + roll_x = np.arctan2(t0, t1) + + # Pitch (y-axis rotation) + t2 = +2.0 * (w * y - z * x) + t2 = np.clip(t2, -1.0, 1.0) + pitch_y = np.arcsin(t2) + + # Yaw (z-axis rotation) + t3 = +2.0 * (w * z + x * y) + t4 = +1.0 - 2.0 * (y * y + z * z) + yaw_z = np.arctan2(t3, t4) + + # Returns roll, pitch, yaw in a NumPy array in radians + return np.array([roll_x, pitch_y, yaw_z]) + + def euler_to_quaternion(self, root_rot): + roll, pitch, yaw = root_rot[0], root_rot[1], root_rot[2] + cy = np.cos(yaw * 0.5) + sy = np.sin(yaw * 0.5) + cp = np.cos(pitch * 0.5) + sp = np.sin(pitch * 0.5) + cr = np.cos(roll * 0.5) + sr = np.sin(roll * 0.5) + + qw = cy * cp * cr + sy * sp * sr + qx = cy * cp * sr - sy * sp * cr + qy = sy * cp * sr + cy * sp * cr + qz = sy * cp * cr - cy * sp * sr + + return np.array([qx, qy, qz, qw]) + + @property + def observation_dim(self): + """Size of AMP observations.""" + return self.trajectories[0].shape[1] + 1 + + @property + def num_motions(self): + return len(self.trajectory_names) + @staticmethod + def get_root_pos(pose): + return pose[0:3] + + @staticmethod + def get_root_pos_batch(poses): + return poses[:, 0:3] + + @staticmethod + def get_root_rot(pose): + return pose[3:7] + + @staticmethod + def get_root_rot_batch(poses): + return poses[:, 3:7] + + @staticmethod + def get_joint_pose_batch_12dof(poses): + return poses[:, 13:25] + + @staticmethod + def get_tar_toe_pos_local(pose): + return pose[G1_AMPLoader.TAR_TOE_POS_LOCAL_START_IDX:G1_AMPLoader.TAR_TOE_POS_LOCAL_END_IDX] + + @staticmethod + def get_tar_toe_pos_local_batch(poses): + return poses[:, G1_AMPLoader.TAR_TOE_POS_LOCAL_START_IDX:G1_AMPLoader.TAR_TOE_POS_LOCAL_END_IDX] diff --git a/rsl_rl/utils/motion_util.py b/rsl_rl/utils/motion_util.py new file mode 100644 index 0000000..3d49bfa --- /dev/null +++ b/rsl_rl/utils/motion_util.py @@ -0,0 +1,97 @@ +# coding=utf-8 +# Copyright 2020 The Google Research Authors. +# +# 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. + +"""Utility functions for processing motion clips.""" + +import os +import inspect +currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) +parentdir = os.path.dirname(os.path.dirname(currentdir)) +os.sys.path.insert(0, parentdir) + +import numpy as np + +from rsl_rl.utils import pose3d +# from pybullet_utils import transformations + + +def standardize_quaternion(q): + """Returns a quaternion where q.w >= 0 to remove redundancy due to q = -q. + + Args: + q: A quaternion to be standardized. + + Returns: + A quaternion with q.w >= 0. + + """ + if q[-1] < 0: + q = -q + return q + + +def normalize_rotation_angle(theta): + """Returns a rotation angle normalized between [-pi, pi]. + + Args: + theta: angle of rotation (radians). + + Returns: + An angle of rotation normalized between [-pi, pi]. + + """ + norm_theta = theta + if np.abs(norm_theta) > np.pi: + norm_theta = np.fmod(norm_theta, 2 * np.pi) + if norm_theta >= 0: + norm_theta += -2 * np.pi + else: + norm_theta += 2 * np.pi + + return norm_theta + + +def calc_heading(q): + """Returns the heading of a rotation q, specified as a quaternion. + + The heading represents the rotational component of q along the vertical + axis (z axis). + + Args: + q: A quaternion that the heading is to be computed from. + + Returns: + An angle representing the rotation about the z axis. + + """ + ref_dir = np.array([1, 0, 0]) + rot_dir = pose3d.QuaternionRotatePoint(ref_dir, q) + heading = np.arctan2(rot_dir[1], rot_dir[0]) + return heading + + +# def calc_heading_rot(q): +# """Return a quaternion representing the heading rotation of q along the vertical axis (z axis). + +# Args: +# q: A quaternion that the heading is to be computed from. + +# Returns: +# A quaternion representing the rotation about the z axis. + +# """ +# heading = calc_heading(q) +# q_heading = transformations.quaternion_about_axis(heading, [0, 0, 1]) +# return q_heading diff --git a/rsl_rl/utils/neptune_utils.py b/rsl_rl/utils/neptune_utils.py new file mode 100644 index 0000000..3796ec8 --- /dev/null +++ b/rsl_rl/utils/neptune_utils.py @@ -0,0 +1,94 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +from dataclasses import asdict +from torch.utils.tensorboard import SummaryWriter + +try: + import neptune +except ModuleNotFoundError: + raise ModuleNotFoundError("neptune-client is required to log to Neptune.") + + +class NeptuneLogger: + def __init__(self, project, token): + self.run = neptune.init_run(project=project, api_token=token) + + def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): + self.run["runner_cfg"] = runner_cfg + self.run["policy_cfg"] = policy_cfg + self.run["alg_cfg"] = alg_cfg + self.run["env_cfg"] = asdict(env_cfg) + + +class NeptuneSummaryWriter(SummaryWriter): + """Summary writer for Neptune.""" + + def __init__(self, log_dir: str, flush_secs: int, cfg): + super().__init__(log_dir, flush_secs) + + try: + project = cfg["neptune_project"] + except KeyError: + raise KeyError("Please specify neptune_project in the runner config, e.g. legged_gym.") + + try: + token = os.environ["NEPTUNE_API_TOKEN"] + except KeyError: + raise KeyError( + "Neptune api token not found. Please run or add to ~/.bashrc: export NEPTUNE_API_TOKEN=YOUR_API_TOKEN" + ) + + try: + entity = os.environ["NEPTUNE_USERNAME"] + except KeyError: + raise KeyError( + "Neptune username not found. Please run or add to ~/.bashrc: export NEPTUNE_USERNAME=YOUR_USERNAME" + ) + + neptune_project = entity + "/" + project + + self.neptune_logger = NeptuneLogger(neptune_project, token) + + self.name_map = { + "Train/mean_reward/time": "Train/mean_reward_time", + "Train/mean_episode_length/time": "Train/mean_episode_length_time", + } + + run_name = os.path.split(log_dir)[-1] + + self.neptune_logger.run["log_dir"].log(run_name) + + def _map_path(self, path): + if path in self.name_map: + return self.name_map[path] + else: + return path + + def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_style=False): + super().add_scalar( + tag, + scalar_value, + global_step=global_step, + walltime=walltime, + new_style=new_style, + ) + self.neptune_logger.run[self._map_path(tag)].log(scalar_value, step=global_step) + + def stop(self): + self.neptune_logger.run.stop() + + def log_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): + self.neptune_logger.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg) + + def save_model(self, model_path, iter): + self.neptune_logger.run["model/saved_model_" + str(iter)].upload(model_path) + + def save_file(self, path, iter=None): + name = path.rsplit("/", 1)[-1].split(".")[0] + self.neptune_logger.run["git_diff/" + name].upload(path) diff --git a/rsl_rl/utils/pose3d.py b/rsl_rl/utils/pose3d.py new file mode 100644 index 0000000..a4cba1e --- /dev/null +++ b/rsl_rl/utils/pose3d.py @@ -0,0 +1,283 @@ +# coding=utf-8 +# Copyright 2020 The Google Research Authors. +# +# 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. +"""Utilities for 3D pose conversion.""" +import math +import numpy as np + +# from pybullet_utils import transformations + +VECTOR3_0 = np.zeros(3, dtype=np.float64) +VECTOR3_1 = np.ones(3, dtype=np.float64) +VECTOR3_X = np.array([1, 0, 0], dtype=np.float64) +VECTOR3_Y = np.array([0, 1, 0], dtype=np.float64) +VECTOR3_Z = np.array([0, 0, 1], dtype=np.float64) + +# QUATERNION_IDENTITY is the multiplicative identity 1.0 + 0i + 0j + 0k. +# When interpreted as a rotation, it is the identity rotation. +QUATERNION_IDENTITY = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float64) + + +def Vector3RandomNormal(sigma, mu=VECTOR3_0): + """Returns a random 3D vector from a normal distribution. + + Each component is selected independently from a normal distribution. + + Args: + sigma: Scale (or stddev) of distribution for all variables. + mu: Mean of distribution for each variable. + + Returns: + A 3D vector in a numpy array. + """ + + random_v3 = np.random.normal(scale=sigma, size=3) + mu + return random_v3 + + +def Vector3RandomUniform(low=VECTOR3_0, high=VECTOR3_1): + """Returns a 3D vector selected uniformly from the input box. + + Args: + low: The min-value corner of the box. + high: The max-value corner of the box. + + Returns: + A 3D vector in a numpy array. + """ + + random_x = np.random.uniform(low=low[0], high=high[0]) + random_y = np.random.uniform(low=low[1], high=high[1]) + random_z = np.random.uniform(low=low[2], high=high[2]) + return np.array([random_x, random_y, random_z]) + + +def Vector3RandomUnit(): + """Returns a random 3D vector with unit length. + + Generates a 3D vector selected uniformly from the unit sphere. + + Returns: + A normalized 3D vector in a numpy array. + """ + longitude = np.random.uniform(low=-math.pi, high=math.pi) + sin_latitude = np.random.uniform(low=-1.0, high=1.0) + cos_latitude = math.sqrt(1.0 - sin_latitude * sin_latitude) + x = math.cos(longitude) * cos_latitude + y = math.sin(longitude) * cos_latitude + z = sin_latitude + return np.array([x, y, z], dtype=np.float64) + + +def QuaternionNormalize(q): + """Normalizes the quaternion to length 1. + + Divides the quaternion by its magnitude. If the magnitude is too + small, returns the quaternion identity value (1.0). + + Args: + q: A quaternion to be normalized. + + Raises: + ValueError: If input quaternion has length near zero. + + Returns: + A quaternion with magnitude 1 in a numpy array [x, y, z, w]. + + """ + q_norm = np.linalg.norm(q) + if np.isclose(q_norm, 0.0): + raise ValueError( + 'Quaternion may not be zero in QuaternionNormalize: |q| = %f, q = %s' % + (q_norm, q)) + return q / q_norm + + +def QuaternionFromAxisAngle(axis, angle): + """Returns a quaternion that generates the given axis-angle rotation. + + Returns the quaternion: sin(angle/2) * axis + cos(angle/2). + + Args: + axis: Axis of rotation, a 3D vector in a numpy array. + angle: The angle of rotation (radians). + + Raises: + ValueError: If input axis is not a normalizable 3D vector. + + Returns: + A unit quaternion in a numpy array. + + """ + if len(axis) != 3: + raise ValueError('Axis vector should have three components: %s' % axis) + axis_norm = np.linalg.norm(axis) + if np.isclose(axis_norm, 0.0): + raise ValueError('Axis vector may not have zero length: |v| = %f, v = %s' % + (axis_norm, axis)) + half_angle = angle * 0.5 + q = np.zeros(4, dtype=np.float64) + q[0:3] = axis + q[0:3] *= math.sin(half_angle) / axis_norm + q[3] = math.cos(half_angle) + return q + + +def QuaternionToAxisAngle(quat, default_axis=VECTOR3_Z, direction_axis=None): + """Calculates axis and angle of rotation performed by a quaternion. + + Calculates the axis and angle of the rotation performed by the quaternion. + The quaternion should have four values and be normalized. + + Args: + quat: Unit quaternion in a numpy array. + default_axis: 3D vector axis used if the rotation is near to zero. Without + this default, small rotations would result in an exception. It is + reasonable to use a default axis for tiny rotations, because zero angle + rotations about any axis are equivalent. + direction_axis: Used to disambiguate rotation directions. If the + direction_axis is specified, the axis of the rotation will be chosen such + that its inner product with the direction_axis is non-negative. + + Raises: + ValueError: If quat is not a normalized quaternion. + + Returns: + axis: Axis of rotation. + angle: Angle in radians. + """ + if len(quat) != 4: + raise ValueError( + 'Quaternion should have four components [x, y, z, w]: %s' % quat) + if not np.isclose(1.0, np.linalg.norm(quat)): + raise ValueError('Quaternion should have unit length: |q| = %f, q = %s' % + (np.linalg.norm(quat), quat)) + axis = quat[:3].copy() + axis_norm = np.linalg.norm(axis) + min_axis_norm = 1e-8 + if axis_norm < min_axis_norm: + axis = default_axis + if len(default_axis) != 3: + raise ValueError('Axis vector should have three components: %s' % axis) + if not np.isclose(np.linalg.norm(axis), 1.0): + raise ValueError('Axis vector should have unit length: |v| = %f, v = %s' % + (np.linalg.norm(axis), axis)) + else: + axis /= axis_norm + sin_half_angle = axis_norm + if direction_axis is not None and np.inner(axis, direction_axis) < 0: + sin_half_angle = -sin_half_angle + axis = -axis + cos_half_angle = quat[3] + half_angle = math.atan2(sin_half_angle, cos_half_angle) + angle = half_angle * 2 + return axis, angle + + +def QuaternionRandomRotation(max_angle=math.pi): + """Creates a random small rotation around a random axis. + + Generates a small rotation with the axis vector selected uniformly + from the unit sphere and an angle selected from a uniform + distribution over [0, max_angle]. + + If the max_angle is not specified, the rotation should be selected + uniformly over all possible rotation angles. + + Args: + max_angle: The maximum angle of rotation (radians). + + Returns: + A unit quaternion in a numpy array. + + """ + + angle = np.random.uniform(low=0, high=max_angle) + axis = Vector3RandomUnit() + return QuaternionFromAxisAngle(axis, angle) + + +# def QuaternionRotatePoint(point, quat): +# """Performs a rotation by quaternion. + +# Rotate the point by the quaternion using quaternion multiplication, +# (q * p * q^-1), without constructing the rotation matrix. + +# Args: +# point: The point to be rotated. +# quat: The rotation represented as a quaternion [x, y, z, w]. + +# Returns: +# A 3D vector in a numpy array. +# """ + +# q_point = np.array([point[0], point[1], point[2], 0.0]) +# quat_inverse = transformations.quaternion_inverse(quat) +# q_point_rotated = transformations.quaternion_multiply( +# transformations.quaternion_multiply(quat, q_point), quat_inverse) +# return q_point_rotated[:3] + + +def IsRotationMatrix(m): + """Returns true if the 3x3 submatrix represents a rotation. + + Args: + m: A transformation matrix. + + Raises: + ValueError: If input is not a matrix of size at least 3x3. + + Returns: + True if the 3x3 submatrix is a rotation (orthogonal). + """ + if len(m.shape) != 2 or m.shape[0] < 3 or m.shape[1] < 3: + raise ValueError('Matrix should be 3x3 or 4x4: %s\n %s' % (m.shape, m)) + rot = m[:3, :3] + eye = np.matmul(rot, np.transpose(rot)) + return np.isclose(eye, np.identity(3), atol=1e-4).all() + +# def ZAxisAlignedRobotPoseTool(robot_pose_tool): +# """Returns the current gripper pose rotated for alignment with the z-axis. + +# Args: +# robot_pose_tool: a pose3d.Pose3d() instance. + +# Returns: +# An instance of pose.Transform representing the current gripper pose +# rotated for alignment with the z-axis. +# """ +# # Align the current pose to the z-axis. +# robot_pose_tool.quaternion = transformations.quaternion_multiply( +# RotationBetween( +# robot_pose_tool.matrix4x4[0:3, 0:3].dot(np.array([0, 0, 1])), +# np.array([0.0, 0.0, -1.0])), robot_pose_tool.quaternion) +# return robot_pose_tool + +# def RotationBetween(a_translation_b, a_translation_c): +# """Computes the rotation from one vector to another. + +# The computed rotation has the property that: + +# a_translation_c = a_rotation_b_to_c * a_translation_b + +# Args: +# a_translation_b: vec3, vector to rotate from +# a_translation_c: vec3, vector to rotate to + +# Returns: +# a_rotation_b_to_c: new Orientation +# """ +# rotation = rotation3.Rotation3.rotation_between( +# a_translation_b, a_translation_c, err_msg='RotationBetween') +# return rotation.quaternion.xyzw diff --git a/rsl_rl/utils/utils.py b/rsl_rl/utils/utils.py new file mode 100644 index 0000000..3605622 --- /dev/null +++ b/rsl_rl/utils/utils.py @@ -0,0 +1,360 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import git +import importlib +import os +import pathlib +import torch +import warnings +from tensordict import TensorDict +from typing import Callable +import numpy as np +class RunningMeanStd: + def __init__(self, epsilon: float = 1e-4, shape: Tuple[int, ...] = ()): + """ + Calculates the running mean and std of a data stream + https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm + :param epsilon: helps with arithmetic issues + :param shape: the shape of the data stream's output + """ + self.mean = np.zeros(shape, np.float64) + self.var = np.ones(shape, np.float64) + self.count = epsilon + + def update(self, arr: np.ndarray) -> None: + batch_mean = np.mean(arr, axis=0) + batch_var = np.var(arr, axis=0) + batch_count = arr.shape[0] + self.update_from_moments(batch_mean, batch_var, batch_count) + + def update_from_moments(self, batch_mean: np.ndarray, batch_var: np.ndarray, batch_count: int) -> None: + delta = batch_mean - self.mean + tot_count = self.count + batch_count + + new_mean = self.mean + delta * batch_count / tot_count + m_a = self.var * self.count + m_b = batch_var * batch_count + m_2 = m_a + m_b + np.square(delta) * self.count * batch_count / (self.count + batch_count) + new_var = m_2 / (self.count + batch_count) + + new_count = batch_count + self.count + + self.mean = new_mean + self.var = new_var + self.count = new_count + + +class Normalizer(RunningMeanStd): + def __init__(self, input_dim, epsilon=1e-4, clip_obs=10.0): + super().__init__(shape=input_dim) + self.epsilon = epsilon + self.clip_obs = clip_obs + + def normalize(self, input): + return np.clip((input - self.mean) / np.sqrt(self.var + self.epsilon), -self.clip_obs, self.clip_obs) + + def normalize_torch(self, input, device): + mean_torch = torch.tensor(self.mean, device=device, dtype=torch.float32) + std_torch = torch.sqrt(torch.tensor(self.var + self.epsilon, device=device, dtype=torch.float32)) + return torch.clamp((input - mean_torch) / std_torch, -self.clip_obs, self.clip_obs) + + def update_normalizer(self, rollouts, expert_loader): + policy_data_generator = rollouts.feed_forward_generator_amp(None, mini_batch_size=expert_loader.batch_size) + expert_data_generator = expert_loader.dataset.feed_forward_generator_amp(expert_loader.batch_size) + + for expert_batch, policy_batch in zip(expert_data_generator, policy_data_generator): + self.update(torch.vstack(tuple(policy_batch) + tuple(expert_batch)).cpu().numpy()) + + +def resolve_nn_activation(act_name: str) -> torch.nn.Module: + """Resolves the activation function from the name. + + Args: + act_name: The name of the activation function. + + Returns: + The activation function. + + Raises: + ValueError: If the activation function is not found. + """ + act_dict = { + "elu": torch.nn.ELU(), + "selu": torch.nn.SELU(), + "relu": torch.nn.ReLU(), + "crelu": torch.nn.CELU(), + "lrelu": torch.nn.LeakyReLU(), + "tanh": torch.nn.Tanh(), + "sigmoid": torch.nn.Sigmoid(), + "softplus": torch.nn.Softplus(), + "gelu": torch.nn.GELU(), + "swish": torch.nn.SiLU(), + "mish": torch.nn.Mish(), + "identity": torch.nn.Identity(), + } + + act_name = act_name.lower() + if act_name in act_dict: + return act_dict[act_name] + else: + raise ValueError(f"Invalid activation function '{act_name}'. Valid activations are: {list(act_dict.keys())}") + + +def resolve_optimizer(optimizer_name: str) -> torch.optim.Optimizer: + """Resolves the optimizer from the name. + + Args: + optimizer_name: The name of the optimizer. + + Returns: + The optimizer. + + Raises: + ValueError: If the optimizer is not found. + """ + optimizer_dict = { + "adam": torch.optim.Adam, + "adamw": torch.optim.AdamW, + "sgd": torch.optim.SGD, + "rmsprop": torch.optim.RMSprop, + } + + optimizer_name = optimizer_name.lower() + if optimizer_name in optimizer_dict: + return optimizer_dict[optimizer_name] + else: + raise ValueError(f"Invalid optimizer '{optimizer_name}'. Valid optimizers are: {list(optimizer_dict.keys())}") + + +def split_and_pad_trajectories( + tensor: torch.Tensor | TensorDict, dones: torch.Tensor +) -> tuple[torch.Tensor | TensorDict, torch.Tensor]: + """Splits trajectories at done indices. Then concatenates them and pads with zeros up to the length of the longest + trajectory. Returns masks corresponding to valid parts of the trajectories. + + Example: + Input: [[a1, a2, a3, a4 | a5, a6], + [b1, b2 | b3, b4, b5 | b6]] + + Output:[[a1, a2, a3, a4], | [[True, True, True, True], + [a5, a6, 0, 0], | [True, True, False, False], + [b1, b2, 0, 0], | [True, True, False, False], + [b3, b4, b5, 0], | [True, True, True, False], + [b6, 0, 0, 0]] | [True, False, False, False]] + + Assumes that the input has the following order of dimensions: [time, number of envs, additional dimensions] + """ + + dones = dones.clone() + dones[-1] = 1 + # Permute the buffers to have order (num_envs, num_transitions_per_env, ...), for correct reshaping + flat_dones = dones.transpose(1, 0).reshape(-1, 1) + # Get length of trajectory by counting the number of successive not done elements + done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero()[:, 0])) + trajectory_lengths = done_indices[1:] - done_indices[:-1] + trajectory_lengths_list = trajectory_lengths.tolist() + # Extract the individual trajectories + if isinstance(tensor, TensorDict): + padded_trajectories = {} + for k, v in tensor.items(): + # split the tensor into trajectories + trajectories = torch.split(v.transpose(1, 0).flatten(0, 1), trajectory_lengths_list) + # add at least one full length trajectory + trajectories = trajectories + (torch.zeros(v.shape[0], *v.shape[2:], device=v.device),) + # pad the trajectories to the length of the longest trajectory + padded_trajectories[k] = torch.nn.utils.rnn.pad_sequence(trajectories) + # remove the added tensor + padded_trajectories[k] = padded_trajectories[k][:, :-1] + padded_trajectories = TensorDict( + padded_trajectories, batch_size=[tensor.batch_size[0], len(trajectory_lengths_list)] + ) + else: + # split the tensor into trajectories + trajectories = torch.split(tensor.transpose(1, 0).flatten(0, 1), trajectory_lengths_list) + # add at least one full length trajectory + trajectories = trajectories + (torch.zeros(tensor.shape[0], *tensor.shape[2:], device=tensor.device),) + # pad the trajectories to the length of the longest trajectory + padded_trajectories = torch.nn.utils.rnn.pad_sequence(trajectories) + # remove the added tensor + padded_trajectories = padded_trajectories[:, :-1] + # create masks for the valid parts of the trajectories + trajectory_masks = trajectory_lengths > torch.arange(0, tensor.shape[0], device=tensor.device).unsqueeze(1) + return padded_trajectories, trajectory_masks + + +def unpad_trajectories(trajectories, masks): + """Does the inverse operation of split_and_pad_trajectories()""" + # Need to transpose before and after the masking to have proper reshaping + return ( + trajectories.transpose(1, 0)[masks.transpose(1, 0)] + .view(-1, trajectories.shape[0], trajectories.shape[-1]) + .transpose(1, 0) + ) + + +def store_code_state(logdir, repositories) -> list: + git_log_dir = os.path.join(logdir, "git") + os.makedirs(git_log_dir, exist_ok=True) + file_paths = [] + for repository_file_path in repositories: + try: + repo = git.Repo(repository_file_path, search_parent_directories=True) + t = repo.head.commit.tree + except Exception: + print(f"Could not find git repository in {repository_file_path}. Skipping.") + # skip if not a git repository + continue + # get the name of the repository + repo_name = pathlib.Path(repo.working_dir).name + diff_file_name = os.path.join(git_log_dir, f"{repo_name}.diff") + # check if the diff file already exists + if os.path.isfile(diff_file_name): + continue + # write the diff file + print(f"Storing git diff for '{repo_name}' in: {diff_file_name}") + with open(diff_file_name, "x", encoding="utf-8") as f: + content = f"--- git status ---\n{repo.git.status()} \n\n\n--- git diff ---\n{repo.git.diff(t)}" + f.write(content) + # add the file path to the list of files to be uploaded + file_paths.append(diff_file_name) + return file_paths + + +def string_to_callable(name: str) -> Callable: + """Resolves the module and function names to return the function. + + Args: + name: The function name. The format should be 'module:attribute_name'. + + Raises: + ValueError: When the resolved attribute is not a function. + ValueError: When unable to resolve the attribute. + + Returns: + The function loaded from the module. + """ + try: + mod_name, attr_name = name.split(":") + mod = importlib.import_module(mod_name) + callable_object = getattr(mod, attr_name) + # check if attribute is callable + if callable(callable_object): + return callable_object + else: + raise ValueError(f"The imported object is not callable: '{name}'") + except AttributeError as e: + msg = ( + "We could not interpret the entry as a callable object. The format of input should be" + f" 'module:attribute_name'\nWhile processing input '{name}', received the error:\n {e}." + ) + raise ValueError(msg) + + +def resolve_obs_groups( + obs: TensorDict, obs_groups: dict[str, list[str]], default_sets: list[str] +) -> dict[str, list[str]]: + """Validates the observation configuration and defaults missing observation sets. + + The input is an observation dictionary `obs` containing observation groups and a configuration dictionary + `obs_groups` where the keys are the observation sets and the values are lists of observation groups. + + The configuration dictionary could for example look like: + { + "policy": ["group_1", "group_2"], + "critic": ["group_1", "group_3"] + } + + This means that the 'policy' observation set will contain the observations "group_1" and "group_2" and the + 'critic' observation set will contain the observations "group_1" and "group_3". This function will check that all + the observations in the 'policy' and 'critic' observation sets are present in the observation dictionary from the + environment. + + Additionally, if one of the `default_sets`, e.g. "critic", is not present in the configuration dictionary, + this function will: + + 1. Check if a group with the same name exists in the observations and assign this group to the observation set. + 2. If 1. fails, it will assign the observations from the 'policy' observation set to the default observation set. + + Args: + obs: Observations from the environment in the form of a dictionary. + obs_groups: Observation sets configuration. + default_sets: Reserved observation set names used by the algorithm (besides 'policy'). + If not provided in 'obs_groups', a default behavior gets triggered. + + Returns: + The resolved observation groups. + + Raises: + ValueError: If any observation set is an empty list. + ValueError: If any observation set contains an observation term that is not present in the observations. + """ + # check if policy observation set exists + if "policy" not in obs_groups.keys(): + if "policy" in obs: + obs_groups["policy"] = ["policy"] + warnings.warn( + "The observation configuration dictionary 'obs_groups' must contain the 'policy' key." + " As an observation group with the name 'policy' was found, this is assumed to be the observation set." + " Consider adding the 'policy' key to the 'obs_groups' dictionary for clarity." + " This behavior will be removed in a future version." + ) + else: + raise ValueError( + "The observation configuration dictionary 'obs_groups' must contain the 'policy' key." + f" Found keys: {list(obs_groups.keys())}" + ) + + # check all observation sets for valid observation groups + for set_name, groups in obs_groups.items(): + # check if the list is empty + if len(groups) == 0: + msg = f"The '{set_name}' key in the 'obs_groups' dictionary can not be an empty list." + if set_name in default_sets: + if set_name not in obs: + msg += " Consider removing the key to default to the observations used for the 'policy' set." + else: + msg += ( + f" Consider removing the key to default to the observation '{set_name}' from the environment." + ) + raise ValueError(msg) + # check groups exist inside the observations from the environment + for group in groups: + if group not in obs: + raise ValueError( + f"Observation '{group}' in observation set '{set_name}' not found in the observations from the" + f" environment. Available observations from the environment: {list(obs.keys())}" + ) + + # fill missing observation sets + for default_set_name in default_sets: + if default_set_name not in obs_groups.keys(): + if default_set_name in obs: + obs_groups[default_set_name] = [default_set_name] + warnings.warn( + f"The observation configuration dictionary 'obs_groups' must contain the '{default_set_name}' key." + f" As an observation group with the name '{default_set_name}' was found, this is assumed to be the" + f" observation set. Consider adding the '{default_set_name}' key to the 'obs_groups' dictionary for" + " clarity. This behavior will be removed in a future version." + ) + else: + obs_groups[default_set_name] = obs_groups["policy"].copy() + warnings.warn( + f"The observation configuration dictionary 'obs_groups' must contain the '{default_set_name}' key." + f" As the configuration for '{default_set_name}' is missing, the observations from the 'policy' set" + f" are used. Consider adding the '{default_set_name}' key to the 'obs_groups' dictionary for" + " clarity. This behavior will be removed in a future version." + ) + + # print the final parsed observation sets + print("-" * 80) + print("Resolved observation sets: ") + for set_name, groups in obs_groups.items(): + print("\t", set_name, ": ", groups) + print("-" * 80) + + return obs_groups diff --git a/rsl_rl/utils/wandb_utils.py b/rsl_rl/utils/wandb_utils.py new file mode 100644 index 0000000..243e82d --- /dev/null +++ b/rsl_rl/utils/wandb_utils.py @@ -0,0 +1,87 @@ +# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import os +from dataclasses import asdict +from torch.utils.tensorboard import SummaryWriter + +try: + import wandb +except ModuleNotFoundError: + raise ModuleNotFoundError("Wandb is required to log to Weights and Biases.") + + +class WandbSummaryWriter(SummaryWriter): + """Summary writer for Weights and Biases.""" + + def __init__(self, log_dir: str, flush_secs: int, cfg): + super().__init__(log_dir, flush_secs) + + # Get the run name + run_name = os.path.split(log_dir)[-1] + + try: + project = cfg["wandb_project"] + except KeyError: + raise KeyError("Please specify wandb_project in the runner config, e.g. legged_gym.") + + try: + entity = os.environ["WANDB_USERNAME"] + except KeyError: + entity = None + + # Initialize wandb + wandb.init(project=project, entity=entity, name=run_name) + + # Add log directory to wandb + wandb.config.update({"log_dir": log_dir}) + + self.name_map = { + "Train/mean_reward/time": "Train/mean_reward_time", + "Train/mean_episode_length/time": "Train/mean_episode_length_time", + } + + def store_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): + wandb.config.update({"runner_cfg": runner_cfg}) + wandb.config.update({"policy_cfg": policy_cfg}) + wandb.config.update({"alg_cfg": alg_cfg}) + try: + wandb.config.update({"env_cfg": env_cfg.to_dict()}) + except Exception: + wandb.config.update({"env_cfg": asdict(env_cfg)}) + + def add_scalar(self, tag, scalar_value, global_step=None, walltime=None, new_style=False): + super().add_scalar( + tag, + scalar_value, + global_step=global_step, + walltime=walltime, + new_style=new_style, + ) + wandb.log({self._map_path(tag): scalar_value}, step=global_step) + + def stop(self): + wandb.finish() + + def log_config(self, env_cfg, runner_cfg, alg_cfg, policy_cfg): + self.store_config(env_cfg, runner_cfg, alg_cfg, policy_cfg) + + def save_model(self, model_path, iter): + wandb.save(model_path, base_path=os.path.dirname(model_path)) + + def save_file(self, path, iter=None): + wandb.save(path, base_path=os.path.dirname(path)) + + """ + Private methods. + """ + + def _map_path(self, path): + if path in self.name_map: + return self.name_map[path] + else: + return path diff --git a/rsl_rl_lib.egg-info/PKG-INFO b/rsl_rl_lib.egg-info/PKG-INFO new file mode 100644 index 0000000..9855570 --- /dev/null +++ b/rsl_rl_lib.egg-info/PKG-INFO @@ -0,0 +1,12 @@ +Metadata-Version: 2.4 +Name: rsl-rl-lib +Version: 3.1.0 +Summary: Fast and simple reinforcement learning algorithms (custom build) +License: BSD-3-Clause +Requires-Python: >=3.9 +Requires-Dist: gitpython +Requires-Dist: numpy +Requires-Dist: onnx +Requires-Dist: tensordict +Requires-Dist: torch +Requires-Dist: torchvision diff --git a/rsl_rl_lib.egg-info/SOURCES.txt b/rsl_rl_lib.egg-info/SOURCES.txt new file mode 100644 index 0000000..b335c3f --- /dev/null +++ b/rsl_rl_lib.egg-info/SOURCES.txt @@ -0,0 +1,105 @@ +pyproject.toml +../rsl_rl/__init__.py +../rsl_rl/algorithms/__init__.py +../rsl_rl/algorithms/amp_ppo.py +../rsl_rl/algorithms/distillation.py +../rsl_rl/algorithms/ppo.py +../rsl_rl/build/lib/rsl_rl/__init__.py +../rsl_rl/build/lib/rsl_rl/algorithms/__init__.py +../rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py +../rsl_rl/build/lib/rsl_rl/algorithms/distillation.py +../rsl_rl/build/lib/rsl_rl/algorithms/ppo.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/__init__.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/__init__.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/amp_ppo.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/distillation.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/algorithms/ppo.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/__init__.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/env/vec_env.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/__init__.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/rnd.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/modules/symmetry.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/__init__.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/memory.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/mlp.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/networks/normalization.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/__init__.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/__init__.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/__init__.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/motion_util.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/pose3d.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/utils.py +../rsl_rl/build/lib/rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py +../rsl_rl/build/lib/rsl_rl/env/__init__.py +../rsl_rl/build/lib/rsl_rl/env/vec_env.py +../rsl_rl/build/lib/rsl_rl/modules/__init__.py +../rsl_rl/build/lib/rsl_rl/modules/actor_critic.py +../rsl_rl/build/lib/rsl_rl/modules/actor_critic_recurrent.py +../rsl_rl/build/lib/rsl_rl/modules/discriminator_multi.py +../rsl_rl/build/lib/rsl_rl/modules/rnd.py +../rsl_rl/build/lib/rsl_rl/modules/student_teacher.py +../rsl_rl/build/lib/rsl_rl/modules/student_teacher_recurrent.py +../rsl_rl/build/lib/rsl_rl/modules/symmetry.py +../rsl_rl/build/lib/rsl_rl/networks/__init__.py +../rsl_rl/build/lib/rsl_rl/networks/memory.py +../rsl_rl/build/lib/rsl_rl/networks/mlp.py +../rsl_rl/build/lib/rsl_rl/networks/normalization.py +../rsl_rl/build/lib/rsl_rl/runners/__init__.py +../rsl_rl/build/lib/rsl_rl/runners/amp_on_policy_runner.py +../rsl_rl/build/lib/rsl_rl/runners/distillation_runner.py +../rsl_rl/build/lib/rsl_rl/runners/on_policy_runner.py +../rsl_rl/build/lib/rsl_rl/storage/__init__.py +../rsl_rl/build/lib/rsl_rl/storage/replay_buffer_multi.py +../rsl_rl/build/lib/rsl_rl/storage/rollout_storage.py +../rsl_rl/build/lib/rsl_rl/utils/__init__.py +../rsl_rl/build/lib/rsl_rl/utils/motion_loader_g1.py +../rsl_rl/build/lib/rsl_rl/utils/motion_util.py +../rsl_rl/build/lib/rsl_rl/utils/neptune_utils.py +../rsl_rl/build/lib/rsl_rl/utils/pose3d.py +../rsl_rl/build/lib/rsl_rl/utils/utils.py +../rsl_rl/build/lib/rsl_rl/utils/wandb_utils.py +../rsl_rl/env/__init__.py +../rsl_rl/env/vec_env.py +../rsl_rl/modules/__init__.py +../rsl_rl/modules/actor_critic.py +../rsl_rl/modules/actor_critic_recurrent.py +../rsl_rl/modules/discriminator_multi.py +../rsl_rl/modules/rnd.py +../rsl_rl/modules/student_teacher.py +../rsl_rl/modules/student_teacher_recurrent.py +../rsl_rl/modules/symmetry.py +../rsl_rl/networks/__init__.py +../rsl_rl/networks/memory.py +../rsl_rl/networks/mlp.py +../rsl_rl/networks/normalization.py +../rsl_rl/runners/__init__.py +../rsl_rl/runners/amp_on_policy_runner.py +../rsl_rl/runners/distillation_runner.py +../rsl_rl/runners/on_policy_runner.py +../rsl_rl/storage/__init__.py +../rsl_rl/storage/replay_buffer_multi.py +../rsl_rl/storage/rollout_storage.py +../rsl_rl/utils/__init__.py +../rsl_rl/utils/motion_loader_g1.py +../rsl_rl/utils/motion_util.py +../rsl_rl/utils/neptune_utils.py +../rsl_rl/utils/pose3d.py +../rsl_rl/utils/utils.py +../rsl_rl/utils/wandb_utils.py +../rsl_rl_lib.egg-info/PKG-INFO +../rsl_rl_lib.egg-info/SOURCES.txt +../rsl_rl_lib.egg-info/dependency_links.txt +../rsl_rl_lib.egg-info/requires.txt +../rsl_rl_lib.egg-info/top_level.txt \ No newline at end of file diff --git a/rsl_rl_lib.egg-info/dependency_links.txt b/rsl_rl_lib.egg-info/dependency_links.txt new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/rsl_rl_lib.egg-info/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/rsl_rl_lib.egg-info/requires.txt b/rsl_rl_lib.egg-info/requires.txt new file mode 100644 index 0000000..810220b --- /dev/null +++ b/rsl_rl_lib.egg-info/requires.txt @@ -0,0 +1,6 @@ +gitpython +numpy +onnx +tensordict +torch +torchvision diff --git a/rsl_rl_lib.egg-info/top_level.txt b/rsl_rl_lib.egg-info/top_level.txt new file mode 100644 index 0000000..84b4467 --- /dev/null +++ b/rsl_rl_lib.egg-info/top_level.txt @@ -0,0 +1 @@ +rsl_rl diff --git a/src/.DS_Store b/src/.DS_Store new file mode 100644 index 0000000..ea46c9b Binary files /dev/null and b/src/.DS_Store differ diff --git a/src/mjlab_husky/.DS_Store b/src/mjlab_husky/.DS_Store new file mode 100644 index 0000000..d59fd6e Binary files /dev/null and b/src/mjlab_husky/.DS_Store differ diff --git a/src/mjlab_husky/__init__.py b/src/mjlab_husky/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mjlab_husky/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..1110013 Binary files /dev/null and b/src/mjlab_husky/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/mjlab_husky/__pycache__/__init__.cpython-313.pyc b/src/mjlab_husky/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..86014e7 Binary files /dev/null and b/src/mjlab_husky/__pycache__/__init__.cpython-313.pyc differ diff --git a/src/mjlab_husky/__pycache__/lerobot_numpy.cpython-312.pyc b/src/mjlab_husky/__pycache__/lerobot_numpy.cpython-312.pyc new file mode 100644 index 0000000..750631d Binary files /dev/null and b/src/mjlab_husky/__pycache__/lerobot_numpy.cpython-312.pyc differ diff --git a/src/mjlab_husky/__pycache__/mujoco_gl.cpython-312.pyc b/src/mjlab_husky/__pycache__/mujoco_gl.cpython-312.pyc new file mode 100644 index 0000000..d44c80c Binary files /dev/null and b/src/mjlab_husky/__pycache__/mujoco_gl.cpython-312.pyc differ diff --git a/src/mjlab_husky/__pycache__/mujoco_gl.cpython-313.pyc b/src/mjlab_husky/__pycache__/mujoco_gl.cpython-313.pyc new file mode 100644 index 0000000..1ac7fda Binary files /dev/null and b/src/mjlab_husky/__pycache__/mujoco_gl.cpython-313.pyc differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/__init__.py b/src/mjlab_husky/asset_zoo/robots/skateboard/__init__.py new file mode 100644 index 0000000..dfa39a3 --- /dev/null +++ b/src/mjlab_husky/asset_zoo/robots/skateboard/__init__.py @@ -0,0 +1 @@ +"""Unitree G1 humanoid skater.""" diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/asset_zoo/robots/skateboard/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..c9fae99 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/__pycache__/g1_skater_constants.cpython-312.pyc b/src/mjlab_husky/asset_zoo/robots/skateboard/__pycache__/g1_skater_constants.cpython-312.pyc new file mode 100644 index 0000000..7b47ca4 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/__pycache__/g1_skater_constants.cpython-312.pyc differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/g1_skater_constants.py b/src/mjlab_husky/asset_zoo/robots/skateboard/g1_skater_constants.py new file mode 100644 index 0000000..1ccac86 --- /dev/null +++ b/src/mjlab_husky/asset_zoo/robots/skateboard/g1_skater_constants.py @@ -0,0 +1,342 @@ +"""Unitree G1 constants.""" + +from pathlib import Path + +import mujoco +import os +from mjlab import MJLAB_SRC_PATH +from mjlab.actuator import BuiltinPositionActuatorCfg +from mjlab.entity import EntityArticulationInfoCfg, EntityCfg +from mjlab.utils.actuator import ( + ElectricActuator, + reflected_inertia_from_two_stage_planetary, +) +from mjlab.utils.os import update_assets +from mjlab.utils.spec_config import CollisionCfg +## +# MJCF and assets. +## + +G1_XML: Path = Path( + os.path.join(os.path.dirname(__file__), "xmls", "g1.xml") +) +assert G1_XML.exists() + +SKATEBOARD_XML: Path = Path( + os.path.join(os.path.dirname(__file__), "xmls", "skateboard.xml") +) +assert SKATEBOARD_XML.exists() + + +def get_assets(meshdir: str) -> dict[str, bytes]: + assets: dict[str, bytes] = {} + update_assets(assets, G1_XML.parent / "assets", meshdir) + return assets + + +def get_g1_spec() -> mujoco.MjSpec: + spec = mujoco.MjSpec.from_file(str(G1_XML)) + spec.assets = get_assets(spec.meshdir) + return spec + +def get_skateboard_spec() -> mujoco.MjSpec: + spec = mujoco.MjSpec.from_file(str(SKATEBOARD_XML)) + spec.assets = get_assets(spec.meshdir) + return spec + +## +# Actuator config. +## + +# Motor specs (from Unitree). +ROTOR_INERTIAS_5020 = ( + 0.139e-4, + 0.017e-4, + 0.169e-4, +) +GEARS_5020 = ( + 1, + 1 + (46 / 18), + 1 + (56 / 16), +) +ARMATURE_5020 = reflected_inertia_from_two_stage_planetary( + ROTOR_INERTIAS_5020, GEARS_5020 +) + +ROTOR_INERTIAS_7520_14 = ( + 0.489e-4, + 0.098e-4, + 0.533e-4, +) +GEARS_7520_14 = ( + 1, + 4.5, + 1 + (48 / 22), +) +ARMATURE_7520_14 = reflected_inertia_from_two_stage_planetary( + ROTOR_INERTIAS_7520_14, GEARS_7520_14 +) + +ROTOR_INERTIAS_7520_22 = ( + 0.489e-4, + 0.109e-4, + 0.738e-4, +) +GEARS_7520_22 = ( + 1, + 4.5, + 5, +) +ARMATURE_7520_22 = reflected_inertia_from_two_stage_planetary( + ROTOR_INERTIAS_7520_22, GEARS_7520_22 +) + +ROTOR_INERTIAS_4010 = ( + 0.068e-4, + 0.0, + 0.0, +) +GEARS_4010 = ( + 1, + 5, + 5, +) +ARMATURE_4010 = reflected_inertia_from_two_stage_planetary( + ROTOR_INERTIAS_4010, GEARS_4010 +) + +ACTUATOR_5020 = ElectricActuator( + reflected_inertia=ARMATURE_5020, + velocity_limit=37.0, + effort_limit=25.0, +) +ACTUATOR_7520_14 = ElectricActuator( + reflected_inertia=ARMATURE_7520_14, + velocity_limit=32.0, + effort_limit=88.0, +) +ACTUATOR_7520_22 = ElectricActuator( + reflected_inertia=ARMATURE_7520_22, + velocity_limit=20.0, + effort_limit=139.0, +) +ACTUATOR_4010 = ElectricActuator( + reflected_inertia=ARMATURE_4010, + velocity_limit=22.0, + effort_limit=5.0, +) + +NATURAL_FREQ = 10 * 2.0 * 3.1415926535 # 10Hz +DAMPING_RATIO = 2.0 + +STIFFNESS_5020 = ARMATURE_5020 * NATURAL_FREQ**2 +STIFFNESS_7520_14 = ARMATURE_7520_14 * NATURAL_FREQ**2 +STIFFNESS_7520_22 = ARMATURE_7520_22 * NATURAL_FREQ**2 +STIFFNESS_4010 = ARMATURE_4010 * NATURAL_FREQ**2 + +DAMPING_5020 = 2.0 * DAMPING_RATIO * ARMATURE_5020 * NATURAL_FREQ +DAMPING_7520_14 = 2.0 * DAMPING_RATIO * ARMATURE_7520_14 * NATURAL_FREQ +DAMPING_7520_22 = 2.0 * DAMPING_RATIO * ARMATURE_7520_22 * NATURAL_FREQ +DAMPING_4010 = 2.0 * DAMPING_RATIO * ARMATURE_4010 * NATURAL_FREQ + +G1_ACTUATOR_5020 = BuiltinPositionActuatorCfg( + target_names_expr=( + ".*_elbow_joint", + ".*_shoulder_pitch_joint", + ".*_shoulder_roll_joint", + ".*_shoulder_yaw_joint", + ), + stiffness=STIFFNESS_5020, + damping=DAMPING_5020, + effort_limit=ACTUATOR_5020.effort_limit, + armature=ACTUATOR_5020.reflected_inertia, +) +G1_ACTUATOR_7520_14 = BuiltinPositionActuatorCfg( + target_names_expr=(".*_hip_pitch_joint", ".*_hip_yaw_joint", "waist_yaw_joint"), + stiffness=STIFFNESS_7520_14, + damping=DAMPING_7520_14, + effort_limit=ACTUATOR_7520_14.effort_limit, + armature=ACTUATOR_7520_14.reflected_inertia, +) +G1_ACTUATOR_7520_22 = BuiltinPositionActuatorCfg( + target_names_expr=(".*_hip_roll_joint", ".*_knee_joint"), + stiffness=STIFFNESS_7520_22, + damping=DAMPING_7520_22, + effort_limit=ACTUATOR_7520_22.effort_limit, + armature=ACTUATOR_7520_22.reflected_inertia, +) + +G1_ACTUATOR_WAIST = BuiltinPositionActuatorCfg( + target_names_expr=("waist_pitch_joint", "waist_roll_joint"), + stiffness=STIFFNESS_5020 * 2, + damping=DAMPING_5020 * 2, + effort_limit=ACTUATOR_5020.effort_limit * 2, + armature=ACTUATOR_5020.reflected_inertia * 2, +) +G1_ACTUATOR_ANKLE = BuiltinPositionActuatorCfg( + target_names_expr=(".*_ankle_pitch_joint", ".*_ankle_roll_joint"), + stiffness=STIFFNESS_5020 * 2, + damping=DAMPING_5020 * 2, + effort_limit=ACTUATOR_5020.effort_limit * 2, + armature=ACTUATOR_5020.reflected_inertia * 2, +) + +SKATEBOARD_ACTUATOR_TRUCK_ROLL = BuiltinPositionActuatorCfg( + target_names_expr=( + ".*_truck_roll_joint", + ), + stiffness=10.0, + damping=0.5, + effort_limit=20.0, +) + +SKATEBOARD_ACTUATOR_WHEEL = BuiltinPositionActuatorCfg( + target_names_expr=( + ".*_wheel_joint", + ), + stiffness=0.0, + damping=0.0, + effort_limit=3.0, +) + +SKATEBOARD_ACTUATOR_TILT = BuiltinPositionActuatorCfg( + target_names_expr=( + "board_tilt_joint", + ), + stiffness=15.0, + damping=0.5, +) +## +# Keyframe config. +## + +PUSH_INIT_KEYFRAME = EntityCfg.InitialStateCfg( + pos=(-0.03, 0.1, 0.78), + joint_pos={ + "left_knee_joint": 0.23, + "left_ankle_pitch_joint": -0.20, + + "right_hip_pitch_joint": -0.7, + "right_knee_joint": 1.17, + "right_ankle_pitch_joint": -0.45, + + "left_shoulder_pitch_joint": -0.03, + "left_shoulder_roll_joint": 0.45, + "left_shoulder_yaw_joint": -0.21, + "left_elbow_joint": 1.32, + + "right_shoulder_pitch_joint": -0.7, + "right_shoulder_roll_joint": -0.845, + "right_shoulder_yaw_joint": 0.83, + "right_elbow_joint": 1.19, + }, + joint_vel={".*": 0.0}, +) + +SKATEBOARD_INIT_KEYFRAME = EntityCfg.InitialStateCfg( + pos=(0, 0, 0.10), + joint_pos={".*": 0.0}, + joint_vel={".*": 0.0}, +) + + +## +# Collision config. +## + +FULL_COLLISION = CollisionCfg( + geom_names_expr=(".*_collision",), + condim={r"^(left|right)_foot[1-7]_collision$": 3, ".*_collision": 1}, + priority={r"^(left|right)_foot[1-7]_collision$": 1}, + friction={r"^(left|right)_foot[1-7]_collision$": (1,)}, +) + +FULL_COLLISION_SKATEBOARD = CollisionCfg( + geom_names_expr=(".*_collision", ), + condim={".*_wheel_collision": 6}, + priority={".*_collision": 1}, + friction={".*_wheel_collision": (1.0, 0.01, 0.001)}, +) + +## +# Final config. +## + + +G1_ARTICULATION = EntityArticulationInfoCfg( + actuators=( + G1_ACTUATOR_5020, + G1_ACTUATOR_7520_14, + G1_ACTUATOR_7520_22, + G1_ACTUATOR_WAIST, + G1_ACTUATOR_ANKLE, + ), + soft_joint_pos_limit_factor=0.9, +) + +SKATEBOARD_ARTICULATION = EntityArticulationInfoCfg( + actuators=( + SKATEBOARD_ACTUATOR_TILT, + SKATEBOARD_ACTUATOR_TRUCK_ROLL, + SKATEBOARD_ACTUATOR_WHEEL, + ), +) + +def get_g1_23dof_robot_cfg() -> EntityCfg: + """Get a fresh G1 robot configuration instance. + + Returns a new EntityCfg instance each time to avoid mutation issues when + the config is shared across multiple places. + """ + return EntityCfg( + init_state=PUSH_INIT_KEYFRAME, + collisions=(FULL_COLLISION,), + spec_fn=get_g1_spec, + articulation=G1_ARTICULATION, + ) + +def get_skateboard_cfg() -> EntityCfg: + """Get a fresh skateboard configuration instance. + + Returns a new EntityCfg instance each time to avoid mutation issues when + the config is shared across multiple places. + """ + return EntityCfg( + init_state=SKATEBOARD_INIT_KEYFRAME, + collisions=(FULL_COLLISION_SKATEBOARD,), + spec_fn=get_skateboard_spec, + articulation=SKATEBOARD_ARTICULATION, + ) + + +G1_23Dof_ACTION_SCALE: dict[str, float] = {} +for a in G1_ARTICULATION.actuators: + assert isinstance(a, BuiltinPositionActuatorCfg) + e = a.effort_limit + s = a.stiffness + names = a.target_names_expr + assert e is not None + for n in names: + G1_23Dof_ACTION_SCALE[n] = 0.25 * e / s + + +if __name__ == "__main__": + import mujoco.viewer as viewer + + from mjlab.entity.entity import Entity + from mjlab.scene import SceneCfg, Scene + from mjlab.terrains import TerrainImporterCfg + + robot = Entity(get_g1_23dof_robot_cfg()) + viewer.launch(robot.spec.compile()) + + skateboard = Entity(get_skateboard_cfg()) + viewer.launch(skateboard.spec.compile()) + + SCENE_CFG = SceneCfg( + terrain=TerrainImporterCfg(terrain_type="plane"), + entities={"robot": get_g1_23dof_robot_cfg(), "skateboard": get_skateboard_cfg()}, + ) + + scene = Scene(SCENE_CFG, device="cuda:0") + scene.to_zip(Path("scene.zip")) diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/head_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/head_link.STL new file mode 100644 index 0000000..2ee5fba Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/head_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_ankle_pitch_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_ankle_pitch_link.STL new file mode 100644 index 0000000..69de849 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_ankle_pitch_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_ankle_roll_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_ankle_roll_link.STL new file mode 100644 index 0000000..8864e9f Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_ankle_roll_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_elbow_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_elbow_link.STL new file mode 100644 index 0000000..1a96d99 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_elbow_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_elbow_link_merge.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_elbow_link_merge.STL new file mode 100644 index 0000000..c3b9dad Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_elbow_link_merge.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_index_0_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_index_0_link.STL new file mode 100644 index 0000000..8069369 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_index_0_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_index_1_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_index_1_link.STL new file mode 100644 index 0000000..89d231d Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_index_1_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_middle_0_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_middle_0_link.STL new file mode 100644 index 0000000..8069369 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_middle_0_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_middle_1_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_middle_1_link.STL new file mode 100644 index 0000000..89d231d Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_middle_1_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_palm_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_palm_link.STL new file mode 100644 index 0000000..7d595ed Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_palm_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_thumb_0_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_thumb_0_link.STL new file mode 100644 index 0000000..3028bb4 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_thumb_0_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_thumb_1_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_thumb_1_link.STL new file mode 100644 index 0000000..d1c080c Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_thumb_1_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_thumb_2_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_thumb_2_link.STL new file mode 100644 index 0000000..8b32e96 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hand_thumb_2_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hip_pitch_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hip_pitch_link.STL new file mode 100644 index 0000000..5b751c7 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hip_pitch_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hip_roll_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hip_roll_link.STL new file mode 100644 index 0000000..778437f Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hip_roll_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hip_yaw_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hip_yaw_link.STL new file mode 100644 index 0000000..383093a Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_hip_yaw_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_knee_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_knee_link.STL new file mode 100644 index 0000000..f2e98e5 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_knee_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_rubber_hand.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_rubber_hand.STL new file mode 100644 index 0000000..c44830f Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_rubber_hand.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_shoulder_pitch_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_shoulder_pitch_link.STL new file mode 100644 index 0000000..e698311 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_shoulder_pitch_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_shoulder_roll_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_shoulder_roll_link.STL new file mode 100644 index 0000000..80bca84 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_shoulder_roll_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_shoulder_yaw_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_shoulder_yaw_link.STL new file mode 100644 index 0000000..281e699 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_shoulder_yaw_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_wrist_pitch_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_wrist_pitch_link.STL new file mode 100644 index 0000000..82cc224 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_wrist_pitch_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_wrist_roll_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_wrist_roll_link.STL new file mode 100644 index 0000000..f3c263a Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_wrist_roll_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_wrist_roll_rubber_hand.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_wrist_roll_rubber_hand.STL new file mode 100644 index 0000000..8fa435b Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_wrist_roll_rubber_hand.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_wrist_yaw_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_wrist_yaw_link.STL new file mode 100644 index 0000000..31be4fd Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/left_wrist_yaw_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/logo_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/logo_link.STL new file mode 100644 index 0000000..e979209 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/logo_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/pelvis.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/pelvis.STL new file mode 100644 index 0000000..691a779 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/pelvis.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/pelvis_contour_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/pelvis_contour_link.STL new file mode 100644 index 0000000..4243433 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/pelvis_contour_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_ankle_pitch_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_ankle_pitch_link.STL new file mode 100644 index 0000000..e77d8a2 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_ankle_pitch_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_ankle_roll_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_ankle_roll_link.STL new file mode 100644 index 0000000..d4261dd Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_ankle_roll_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_elbow_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_elbow_link.STL new file mode 100644 index 0000000..f259e38 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_elbow_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_elbow_link_merge.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_elbow_link_merge.STL new file mode 100644 index 0000000..83ce0ba Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_elbow_link_merge.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_index_0_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_index_0_link.STL new file mode 100644 index 0000000..f87ad32 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_index_0_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_index_1_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_index_1_link.STL new file mode 100644 index 0000000..6dea51a Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_index_1_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_middle_0_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_middle_0_link.STL new file mode 100644 index 0000000..f87ad32 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_middle_0_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_middle_1_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_middle_1_link.STL new file mode 100644 index 0000000..6dea51a Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_middle_1_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_palm_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_palm_link.STL new file mode 100644 index 0000000..5ae00a7 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_palm_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_thumb_0_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_thumb_0_link.STL new file mode 100644 index 0000000..1cae7f1 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_thumb_0_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_thumb_1_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_thumb_1_link.STL new file mode 100644 index 0000000..c141fbf Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_thumb_1_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_thumb_2_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_thumb_2_link.STL new file mode 100644 index 0000000..e942923 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hand_thumb_2_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hip_pitch_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hip_pitch_link.STL new file mode 100644 index 0000000..998a0a0 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hip_pitch_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hip_roll_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hip_roll_link.STL new file mode 100644 index 0000000..47b2eeb Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hip_roll_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hip_yaw_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hip_yaw_link.STL new file mode 100644 index 0000000..3718564 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_hip_yaw_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_knee_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_knee_link.STL new file mode 100644 index 0000000..76d21a3 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_knee_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_rubber_hand.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_rubber_hand.STL new file mode 100644 index 0000000..0aacffb Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_rubber_hand.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_shoulder_pitch_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_shoulder_pitch_link.STL new file mode 100644 index 0000000..3f5b4ed Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_shoulder_pitch_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_shoulder_roll_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_shoulder_roll_link.STL new file mode 100644 index 0000000..179d617 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_shoulder_roll_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_shoulder_yaw_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_shoulder_yaw_link.STL new file mode 100644 index 0000000..2ba6076 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_shoulder_yaw_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_wrist_pitch_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_wrist_pitch_link.STL new file mode 100644 index 0000000..da19454 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_wrist_pitch_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_wrist_roll_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_wrist_roll_link.STL new file mode 100644 index 0000000..26868d2 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_wrist_roll_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_wrist_roll_rubber_hand.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_wrist_roll_rubber_hand.STL new file mode 100644 index 0000000..c365aa9 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_wrist_roll_rubber_hand.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_wrist_yaw_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_wrist_yaw_link.STL new file mode 100644 index 0000000..d788902 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/right_wrist_yaw_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_constraint_L_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_constraint_L_link.STL new file mode 100644 index 0000000..75d82f5 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_constraint_L_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_constraint_L_rod_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_constraint_L_rod_link.STL new file mode 100644 index 0000000..6747f3f Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_constraint_L_rod_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_constraint_R_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_constraint_R_link.STL new file mode 100644 index 0000000..5cb5958 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_constraint_R_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_constraint_R_rod_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_constraint_R_rod_link.STL new file mode 100644 index 0000000..95cf415 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_constraint_R_rod_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_link.STL new file mode 100644 index 0000000..17745af Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_link_23dof_rev_1_0.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_link_23dof_rev_1_0.STL new file mode 100644 index 0000000..079aa81 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_link_23dof_rev_1_0.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_link_rev_1_0.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_link_rev_1_0.STL new file mode 100644 index 0000000..8a759a7 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/torso_link_rev_1_0.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_constraint_L.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_constraint_L.STL new file mode 100644 index 0000000..911410f Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_constraint_L.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_constraint_R.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_constraint_R.STL new file mode 100644 index 0000000..babe794 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_constraint_R.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_roll_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_roll_link.STL new file mode 100644 index 0000000..65831ab Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_roll_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_roll_link_rev_1_0.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_roll_link_rev_1_0.STL new file mode 100644 index 0000000..a64f330 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_roll_link_rev_1_0.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_support_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_support_link.STL new file mode 100644 index 0000000..63660fb Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_support_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_yaw_link.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_yaw_link.STL new file mode 100644 index 0000000..7d36b02 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_yaw_link.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_yaw_link_rev_1_0.STL b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_yaw_link_rev_1_0.STL new file mode 100644 index 0000000..0fabb63 Binary files /dev/null and b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/assets/waist_yaw_link_rev_1_0.STL differ diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/g1.xml b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/g1.xml new file mode 100644 index 0000000..8b2cb6e --- /dev/null +++ b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/g1.xml @@ -0,0 +1,326 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/skateboard.xml b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/skateboard.xml new file mode 100644 index 0000000..94e9e69 --- /dev/null +++ b/src/mjlab_husky/asset_zoo/robots/skateboard/xmls/skateboard.xml @@ -0,0 +1,65 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/src/mjlab_husky/envs/__init__.py b/src/mjlab_husky/envs/__init__.py new file mode 100644 index 0000000..7e4a448 --- /dev/null +++ b/src/mjlab_husky/envs/__init__.py @@ -0,0 +1,6 @@ +from mjlab.envs.manager_based_rl_env import ManagerBasedRlEnv as ManagerBasedRlEnv +from mjlab.envs.manager_based_rl_env import ManagerBasedRlEnvCfg as ManagerBasedRlEnvCfg +from mjlab.envs.types import VecEnvObs as VecEnvObs +from mjlab.envs.types import VecEnvStepReturn as VecEnvStepReturn +from mjlab_husky.envs.g1_skate_rl_env import G1SkaterManagerBasedRlEnvCfg as G1SkaterManagerBasedRlEnvCfg +from mjlab_husky.envs.g1_skate_rl_env import G1SkaterManagerBasedRlEnv as G1SkaterManagerBasedRlEnv diff --git a/src/mjlab_husky/envs/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/envs/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..3a6f0c7 Binary files /dev/null and b/src/mjlab_husky/envs/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/mjlab_husky/envs/__pycache__/g1_skate_rl_env.cpython-312.pyc b/src/mjlab_husky/envs/__pycache__/g1_skate_rl_env.cpython-312.pyc new file mode 100644 index 0000000..c8cc8c5 Binary files /dev/null and b/src/mjlab_husky/envs/__pycache__/g1_skate_rl_env.cpython-312.pyc differ diff --git a/src/mjlab_husky/envs/g1_skate_rl_env.py b/src/mjlab_husky/envs/g1_skate_rl_env.py new file mode 100644 index 0000000..7d4a849 --- /dev/null +++ b/src/mjlab_husky/envs/g1_skate_rl_env.py @@ -0,0 +1,538 @@ +from dataclasses import dataclass, field +import mujoco +import numpy as np +import torch +import warp as wp +from prettytable import PrettyTable +from mjlab.envs import types +from mjlab.envs.manager_based_rl_env import ManagerBasedRlEnv +from mjlab.envs.manager_based_rl_env import ManagerBasedRlEnvCfg +from mjlab.managers.reward_manager import RewardManager, RewardTermCfg +from mjlab.scene import Scene +from mjlab.sim.sim import Simulation +from mjlab.utils.logging import print_info +from mjlab.viewer.offscreen_renderer import OffscreenRenderer +from mjlab.utils.lab_api.math import ( + subtract_frame_transforms, + quat_apply, + quat_mul, + matrix_from_quat, +) +from mjlab.viewer.debug_visualizer import DebugVisualizer +_DESIRED_FRAME_COLORS = ((1.0, 0.5, 0.5), (0.5, 1.0, 0.5), (0.5, 0.5, 1.0)) + +@dataclass(kw_only=True) +class G1SkaterManagerBasedRlEnvCfg(ManagerBasedRlEnvCfg): + + push_rewards: dict[str, RewardTermCfg] = field(default_factory=dict) + steer_rewards: dict[str, RewardTermCfg] = field(default_factory=dict) + transition_rewards: dict[str, RewardTermCfg] = field(default_factory=dict) + regularization_rewards: dict[str, RewardTermCfg] = field(default_factory=dict) + + cycle_time: float = 6.0 + phase_ratios: list[float] = field(default_factory=list) + beizer_names: list[str] = field(default_factory=list) + slerp_names: list[str] = field(default_factory=list) + steer_init_pos: list[float] = field(default_factory=list) + rake_angle: float = 60.0 + eval_mode: bool = False + """Whether in evaluation mode. If True, will save metrics to JSON and exit after all episodes complete.""" + eval_output_dir: str | None = None + """Directory to save eval metrics JSON files. If None, saves to current directory.""" + +class G1SkaterManagerBasedRlEnv(ManagerBasedRlEnv): + """Manager-based RL environment.""" + + is_vector_env = True + metadata = { + "render_modes": [None, "rgb_array"], + "mujoco_version": mujoco.__version__, + "warp_version": wp.config.version, + } + cfg: G1SkaterManagerBasedRlEnvCfg # type: ignore[assignment] + + def __init__( + self, + cfg: G1SkaterManagerBasedRlEnvCfg, + device: str, + render_mode: str | None = None, + **kwargs, + ) -> None: + # Initialize base environment state. + self.cfg = cfg # type: ignore[assignment] + if self.cfg.seed is not None: + self.cfg.seed = self.seed(self.cfg.seed) + self._sim_step_counter = 0 + self.extras = {} + self.obs_buf = {} + + # Initialize scene and simulation. + self.scene = Scene(self.cfg.scene, device=device) + self.sim = Simulation( + num_envs=self.scene.num_envs, + cfg=self.cfg.sim, + model=self.scene.compile(), + device=device, + ) + + self.scene.initialize( + mj_model=self.sim.mj_model, + model=self.sim.model, + data=self.sim.data, + ) + + # Print environment info. + print_info("") + table = PrettyTable() + table.title = "Base Environment" + table.field_names = ["Property", "Value"] + table.align["Property"] = "l" + table.align["Value"] = "l" + table.add_row(["Number of environments", self.num_envs]) + table.add_row(["Environment device", self.device]) + table.add_row(["Environment seed", self.cfg.seed]) + table.add_row(["Physics step-size", self.physics_dt]) + table.add_row(["Environment step-size", self.step_dt]) + print_info(table.get_string()) + print_info("") + + self.cycle_time = self.cfg.cycle_time + self.robot = self.scene["robot"] + self.skateboard = self.scene["skateboard"] + self._init_buffers() + + # Initialize RL-specific state. + self.common_step_counter = 0 + self.episode_length_buf = torch.zeros( + cfg.scene.num_envs, device=device, dtype=torch.long + ) + self.render_mode = render_mode + self._offline_renderer: OffscreenRenderer | None = None + if self.render_mode == "rgb_array": + renderer = OffscreenRenderer( + model=self.sim.mj_model, cfg=self.cfg.viewer, scene=self.scene + ) + renderer.initialize() + self._offline_renderer = renderer + self.metadata["render_fps"] = 1.0 / self.step_dt # type: ignore + + # Load all managers. + self.load_managers() + self.setup_manager_visualizers() + + def _init_buffers(self): + self._init_ids_buffers() + self.phase_ratios = torch.tensor(self.cfg.phase_ratios, device=self.device).repeat(self.num_envs, 1) + self.steer_init_pos = torch.tensor(self.cfg.steer_init_pos, device=self.device).repeat(self.num_envs, 1) + self.last_contacts = torch.zeros(self.num_envs, 2, dtype=torch.bool, device=self.device, requires_grad=False) + self.last_wheel_contacts = torch.zeros(self.num_envs, 4, dtype=torch.bool, device=self.device, requires_grad=False) + self.last_contacts_b = torch.zeros(self.num_envs, 2, dtype=torch.bool, device=self.device, requires_grad=False) + self.last_contacts_g = torch.zeros(self.num_envs, 2, dtype=torch.bool, device=self.device, requires_grad=False) + self.contact_phase = torch.zeros(self.num_envs, 4, dtype=torch.float, device=self.device, requires_grad=False) + + self.phase_length_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.long, requires_grad=False) + self.still = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device, requires_grad=False) + + push_init_body_pose = torch.from_numpy(np.load("dataset/ref_pose/push_start_pose_b.npy")).to(self.device).repeat(self.num_envs, 1 , 1) + steer_init_body_pose = torch.from_numpy(np.load("dataset/ref_pose/steer_start_pose_b.npy")).to(self.device).repeat(self.num_envs, 1 , 1) + self.push_init_body_pos_b = push_init_body_pose[..., :3] + self.steer_init_body_pos_b = steer_init_body_pose[..., :3] + self.push_init_body_quat_b = push_init_body_pose[..., 3:] + self.steer_init_body_quat_b = steer_init_body_pose[..., 3:] + self.body_bezier_buffers = { + "push2steer_start_pos_b": torch.zeros(self.num_envs,self.robot.num_bodies, 3, device=self.device, requires_grad=False), + "steer2push_start_pos_b": torch.zeros(self.num_envs, self.robot.num_bodies, 3, device=self.device, requires_grad=False), + "push2steer_start_quat_b": torch.zeros(self.num_envs, self.robot.num_bodies, 4, device=self.device, requires_grad=False), + "steer2push_start_quat_b": torch.zeros(self.num_envs, self.robot.num_bodies, 4, device=self.device, requires_grad=False), + } + + def _init_ids_buffers(self): + self.wheel_body_ids, _ = self.skateboard.find_bodies(name_keys=[".*_wheel"], preserve_order=True) + self.feet_body_ids, _ = self.robot.find_bodies(name_keys=["left_ankle_roll_link", "right_ankle_roll_link"], preserve_order=True) + self.marker_body_ids, _ = self.skateboard.find_sites(name_keys=[".*_marker"], preserve_order=True) + self.left_foot_site_ids, _ = self.robot.find_sites(name_keys=["left_foot_1", "left_foot_2", "left_foot_3", "left_foot_4"], preserve_order=True) + self.right_foot_site_ids, _ = self.robot.find_sites(name_keys=["right_foot_1", "right_foot_2", "right_foot_3", "right_foot_4"], preserve_order=True) + self.beizer_ids, _ = self.robot.find_bodies(name_keys=self.cfg.beizer_names, preserve_order=True) + self.slerp_ids, _ = self.robot.find_bodies(name_keys=self.cfg.slerp_names, preserve_order=True) + self.truck_roll_joint_ids, _ = self.skateboard.find_joints(name_keys=["front_truck_roll_joint", "rear_truck_roll_joint"], preserve_order=True) + + def load_managers(self) -> None: + super().load_managers() + + self.push_reward_manager = RewardManager( + self.cfg.push_rewards, self, scale_by_dt=self.cfg.scale_rewards_by_dt + ) + print_info(f"[INFO] {self.push_reward_manager}") + self.steer_reward_manager = RewardManager( + self.cfg.steer_rewards, self, scale_by_dt=self.cfg.scale_rewards_by_dt + ) + print_info(f"[INFO] {self.steer_reward_manager}") + self.transition_reward_manager = RewardManager( + self.cfg.transition_rewards, self, scale_by_dt=self.cfg.scale_rewards_by_dt + ) + print_info(f"[INFO] {self.transition_reward_manager}") + self.reg_reward_manager = RewardManager( + self.cfg.regularization_rewards, self, scale_by_dt=self.cfg.scale_rewards_by_dt + ) + print_info(f"[INFO] {self.reg_reward_manager}") + + def get_heading_target_w(self, command_name: str) -> torch.Tensor | None: + terms = getattr(self.command_manager, "_terms", None) or getattr(self.command_manager, "terms", None) + if terms is None: + return None + term = terms.get(command_name) + if term is None: + return None + return getattr(term, "target_heading_w", None) + + def step(self, action: torch.Tensor) -> types.VecEnvStepReturn: + self.action_manager.process_action(action.to(self.device)) + self.still = self.command_manager.get_command("skate")[:, 0] < 0.1 # pyright: ignore[reportOptionalSubscript] + for _ in range(self.cfg.decimation): + self._sim_step_counter += 1 + self.action_manager.apply_action() + self.scene.write_data_to_sim() + # self._set_skatedboard_joint_pos() + self.sim.step() + self.scene.update(dt=self.physics_dt) + + # Update env counters. + self.episode_length_buf += 1 + self.phase_length_buf += 1 + self.common_step_counter += 1 + self._compute_contact() + + # Check terminations. + self.reset_buf = self.termination_manager.compute() + self.reset_terminated = self.termination_manager.terminated + self.reset_time_outs = self.termination_manager.time_outs + + contact_coef = self.contact_phase.clone() + push_reward_buf = self.push_reward_manager.compute(self.step_dt) * contact_coef[:,0] + steer_reward_buf = self.steer_reward_manager.compute(self.step_dt) * contact_coef[:,1] + reg_reward_buf = self.reg_reward_manager.compute(self.step_dt) + transition_reward_buf = self.transition_reward_manager.compute(self.step_dt) * torch.logical_or(contact_coef[:, 2], contact_coef[:,3]) + + self.reward_buf = steer_reward_buf + push_reward_buf + reg_reward_buf + transition_reward_buf + self.reset_env_ids = self.reset_buf.nonzero(as_tuple=False).squeeze(-1) + if len(self.reset_env_ids) > 0: + self._reset_idx(self.reset_env_ids) + self.scene.write_data_to_sim() + self.sim.forward() + + self.command_manager.compute(dt=self.step_dt) + + if "interval" in self.event_manager.available_modes: + self.event_manager.apply(mode="interval", dt=self.step_dt) + + self.obs_buf = self.observation_manager.compute(update_history=True) + return ( + self.obs_buf, + self.reward_buf, + self.reset_terminated, + self.reset_time_outs, + self.extras, + ) + + def _set_skatedboard_joint_pos(self): + joint_pos = self.skateboard.data.joint_pos + tilt_joint_pos = joint_pos[:, 0] + rake_angle_rad = torch.deg2rad(torch.tensor(self.cfg.rake_angle, device=self.device, dtype=torch.float32)) + truck_pos = -torch.atan(torch.sin(tilt_joint_pos) * torch.tan(rake_angle_rad)) + truck_pos = torch.clip(truck_pos, -0.1, 0.1) + self.skateboard.write_joint_position_to_sim( + truck_pos.unsqueeze(-1), + joint_ids=self.truck_roll_joint_ids + ) + + def update_visualizers(self, visualizer: DebugVisualizer) -> None: + super().update_visualizers(visualizer) + self._visualize_transition_target(visualizer) + self._visualize_contact_phase(visualizer) + + def _reset_idx(self, env_ids: torch.Tensor | None = None) -> None: + super()._reset_idx(env_ids) + + info = self.push_reward_manager.reset(env_ids) + self.extras["log"].update(info) + info = self.steer_reward_manager.reset(env_ids) + self.extras["log"].update(info) + info = self.reg_reward_manager.reset(env_ids) + self.extras["log"].update(info) + info = self.transition_reward_manager.reset(env_ids) + self.extras["log"].update(info) + + self.phase_length_buf[env_ids] = 0 + for buf in self.body_bezier_buffers.values(): + buf[env_ids] = 0 + + def _compute_contact(self): + self.skateboard_contact_sensor = self.scene.sensors["skateboard_collision"] + self.left_feet_contact_ground = self.scene.sensors["left_feet_ground_contact"] + self.right_feet_contact_ground = self.scene.sensors["right_feet_ground_contact"] + l_contact = torch.norm(self.left_feet_contact_ground.data.force, dim=-1) > 2. + r_contact = torch.norm(self.right_feet_contact_ground.data.force, dim=-1) > 2. + contact = torch.cat([l_contact, r_contact], dim=-1).squeeze(1) + + wheel_contact = torch.logical_or(torch.norm(self.skateboard_contact_sensor.data.force, dim=-1) > 1., + torch.abs(self.skateboard.data.body_link_pos_w[:, self.wheel_body_ids, 2] - 0.03) < 0.005) + self.contact_filt = torch.logical_or(contact, self.last_contacts) + self.wheel_contact_filt = torch.logical_or(wheel_contact, self.last_wheel_contacts) + self.last_contacts = contact + self.last_wheel_contacts = wheel_contact + self._resample_contact_phases() + + def _resample_contact_phases(self): + self.last_contact_phase = self.contact_phase.clone() + phase = self._get_phase() + + push_phase = (phase >= self.phase_ratios[:, 0]) & (phase < self.phase_ratios[:, 1]) & ~self.still + push2steer = (phase >= self.phase_ratios[:, 1]) & (phase < self.phase_ratios[:, 2]) & ~self.still + steer_phase = (phase >= self.phase_ratios[:, 2]) & (phase < self.phase_ratios[:, 3]) & ~self.still + steer2push = (phase >= self.phase_ratios[:, 3]) & (phase <= self.phase_ratios[:, 4]) & ~self.still + + self.contact_phase[:,0] = push_phase.float() + self.contact_phase[:,1] = steer_phase.float() + self.contact_phase[:,2] = push2steer.float() + self.contact_phase[:,3] = steer2push.float() + + self.just_entered_push2steer = push2steer & (self.last_contact_phase[:, 2] < 0.5) + self.just_entered_steer2push = steer2push & (self.last_contact_phase[:, 3] < 0.5) + self.just_exited_push2steer = (self.last_contact_phase[:, 2] > 0.5) & ~push2steer + self.just_exited_steer2push = (self.last_contact_phase[:, 3] > 0.5) & ~steer2push + + body_pos_w = self.robot.data.body_link_pos_w + body_quat_w = self.robot.data.body_link_quat_w + root_pos_w = self.skateboard.data.root_link_pos_w[:, None, :].repeat(1, self.robot.num_bodies, 1) + root_quat_w = self.skateboard.data.root_link_quat_w[:, None, :].repeat(1, self.robot.num_bodies, 1) + body_pos_b, body_quat_b = subtract_frame_transforms(root_pos_w, root_quat_w, body_pos_w, body_quat_w) + if self.just_entered_push2steer.any(): + self.body_bezier_buffers["push2steer_start_pos_b"][self.just_entered_push2steer] = body_pos_b[self.just_entered_push2steer] + self.body_bezier_buffers["push2steer_start_quat_b"][self.just_entered_push2steer] = body_quat_b[self.just_entered_push2steer] + + if self.just_entered_steer2push.any(): + self.body_bezier_buffers["steer2push_start_pos_b"][self.just_entered_steer2push] = body_pos_b[self.just_entered_steer2push] + self.body_bezier_buffers["steer2push_start_quat_b"][self.just_entered_steer2push] = body_quat_b[self.just_entered_steer2push] + + def _get_phase(self): + self.phase_length_buf[self.still] = torch.where((self.phase_length_buf[self.still]-1) % int(self.cycle_time/ 2 / self.step_dt) == 0, + 0, + self.phase_length_buf[self.still]) + phase = ((self.phase_length_buf * self.step_dt / self.cycle_time)) % 1.0 + phase = torch.clip(phase, 0.0, 1.0) + + return phase + + def _steer_remaining_steps(self): + phase = self._get_phase() + steer_end_phase = self.phase_ratios[:, 3] + remaining_phase = torch.where( + phase < steer_end_phase, + steer_end_phase - phase, + 1.0 - phase + steer_end_phase + ) + remaining_steps = remaining_phase * self.cycle_time / self.step_dt + + return remaining_steps + + def _get_feet_contact_b(self): + left_contact_sensor = self.scene.sensors["left_feet_board_contact"] + right_contact_sensor = self.scene.sensors["right_feet_board_contact"] + left_contact_b = (torch.norm(left_contact_sensor.data.force, dim=-1) > 5) + right_contact_b = (torch.norm(right_contact_sensor.data.force, dim=-1) > 5) + contact_b = torch.cat([left_contact_b, right_contact_b], dim=-1).view(self.num_envs, 2) + contact_filt = torch.logical_or(contact_b, self.last_contacts_b) + self.last_contacts_b = contact_b + return contact_filt + + def _get_feet_contact_g(self): + left_contact_ground_sensor = self.scene.sensors["left_feet_ground_contact"] + right_contact_ground_sensor = self.scene.sensors["right_feet_ground_contact"] + left_contact_g = (torch.norm(left_contact_ground_sensor.data.force, dim=-1) > 5) + right_contact_g = (torch.norm(right_contact_ground_sensor.data.force, dim=-1) > 5) + contact_g = torch.cat([left_contact_g, right_contact_g], dim=-1).view(self.num_envs, 2) + + contact_filt = torch.logical_or(contact_g, self.last_contacts_g) + self.last_contacts_g = contact_g + return contact_filt + + def _get_feet_marker_dis(self): + feet_pos = self.robot.data.body_link_pos_w[:, self.feet_body_ids, :3] + marker_pos = self.skateboard.data.site_pos_w[:, self.marker_body_ids, :3] + dis = marker_pos - feet_pos + return dis + + def _get_transition_target_b(self): + phase = self._get_phase() + push2steer = (phase > self.phase_ratios[:, 1]) & (phase < self.phase_ratios[:, 2]) & ~self.still + steer2push = (phase > self.phase_ratios[:, 3]) & (phase < self.phase_ratios[:, 4]) & ~self.still + in_transition = push2steer | steer2push + + body_pos_w = self.robot.data.body_link_pos_w + body_quat_w = self.robot.data.body_link_quat_w + skate_pos_w = self.skateboard.data.root_link_pos_w[:, None, :].repeat(1, self.robot.num_bodies, 1) + skate_quat_w = self.skateboard.data.root_link_quat_w[:, None, :].repeat(1, self.robot.num_bodies, 1) + + current_body_pos_b,current_body_quat_b = subtract_frame_transforms(skate_pos_w, skate_quat_w, body_pos_w, body_quat_w) + target_pos_b = current_body_pos_b.clone() + target_quat_b = current_body_quat_b.clone() + if in_transition.any(): + t = torch.zeros(self.num_envs, device=self.device) + + if push2steer.any(): + t[push2steer] = (phase[push2steer] - self.phase_ratios[push2steer, 1]) / 0.1 + t[push2steer] = torch.clamp(t[push2steer], 0.0, 1.0) + + start_pos_b = self.body_bezier_buffers["push2steer_start_pos_b"][push2steer] + start_quat_b = self.body_bezier_buffers["push2steer_start_quat_b"][push2steer] + + end_pos_b = self.steer_init_body_pos_b[push2steer] + end_quat_b = self.steer_init_body_quat_b[push2steer] + + target_pos_b[push2steer] = bezier_curve(start_pos_b, end_pos_b, t[push2steer], offset=0.2) + target_quat_b[push2steer] = quaternion_slerp(start_quat_b, end_quat_b, t[push2steer]) + + if steer2push.any(): + t[steer2push] = (phase[steer2push] - self.phase_ratios[steer2push, 3]) / 0.05 + t[steer2push] = torch.clamp(t[steer2push], 0.0, 1.0) + + start_pos_b = self.body_bezier_buffers["steer2push_start_pos_b"][steer2push] + start_quat_b = self.body_bezier_buffers["steer2push_start_quat_b"][steer2push] + + end_pos_b = self.push_init_body_pos_b[steer2push] + end_quat_b = self.push_init_body_quat_b[steer2push] + + target_pos_b[steer2push] = bezier_curve(start_pos_b, end_pos_b, t[steer2push],offset=0.2) + target_quat_b[steer2push] = quaternion_slerp(start_quat_b, end_quat_b, t[steer2push]) + + return target_pos_b, target_quat_b, in_transition + + def get_amp_observations(self): + return self.robot.data.joint_pos + + def _visualize_transition_target(self, visualizer: DebugVisualizer): + target_pos_b, target_quat_b, in_transition = self._get_transition_target_b() + + if in_transition.any(): + + root_pos_w = self.skateboard.data.root_link_pos_w[:, :3][:,None,:].repeat(1,self.robot.num_bodies,1) + root_quat_w = self.skateboard.data.root_link_quat_w[:,None,:].repeat(1,self.robot.num_bodies,1) + target_pos_w = root_pos_w[in_transition] + quat_apply( + root_quat_w[in_transition], + target_pos_b[in_transition] + ) + target_quat_w = quat_mul(root_quat_w[in_transition], target_quat_b[in_transition]) + desired_body_rotm = matrix_from_quat(target_quat_w).cpu().numpy() + for i, idx in enumerate(self.beizer_ids): + visualizer.add_frame( + position=target_pos_w[visualizer.env_idx,idx], + rotation_matrix=desired_body_rotm[visualizer.env_idx,idx], + scale=0.1, + label=f"desired_{idx}", + axis_colors=_DESIRED_FRAME_COLORS, + ) + for i, idx in enumerate(self.beizer_ids): + visualizer.add_sphere( + center=target_pos_w[visualizer.env_idx,idx], + radius=0.03, + color=(1.0, 1.0, 0.0, 1.0), + label=f"desired_{idx}", + ) + + def _visualize_contact_phase(self, visualizer: DebugVisualizer): + contact_phase = self.contact_phase.clone() + push_phase = contact_phase[:, 0] + steer_phase = contact_phase[:, 1] + transition_phase = torch.logical_or(contact_phase[:, 2], contact_phase[:, 3]) + target_pos_w = self.robot.data.root_link_pos_w + target_pos_w[...,2] += 0.75 + if push_phase.any(): + visualizer.add_sphere( + center=target_pos_w[visualizer.env_idx], + radius=0.05, + color=(1.0, 0.0, 0.0, 1.0), + label="push_phase", + ) + if steer_phase.any(): + visualizer.add_sphere( + center=target_pos_w[visualizer.env_idx], + radius=0.05, + color=(0.0, 1.0, 0.0, 1.0), + label="steer_phase", + ) + if transition_phase.any(): + visualizer.add_sphere( + center=target_pos_w[visualizer.env_idx], + radius=0.05, + color=(1.0, 1.0, 0.0, 1.0), + label="transition_phase", + ) + + +def bezier_curve(start_p, end_p, t, offset=0.15): + + # middle control point, shape (num_envs, bodys, 3) + middle_p = (start_p + end_p) / 2.0 + middle_p[..., 2] += offset + + t = torch.clamp(t, 0.0, 1.0).view(-1, 1, 1) + + result_pos = (1 - t) ** 2 * start_p + 2 * (1 - t) * t * middle_p + t ** 2 * end_p + + if result_pos.shape[1] == 1: + result_pos = result_pos.squeeze(1) + if result_pos.shape[0] == 1: + result_pos = result_pos.squeeze(0) + + return result_pos + +def quaternion_slerp(q0, q1, t, shortestpath=True): + if t.dim() == 0: + t = t.view(1, 1, 1).expand(q0.shape[0], q0.shape[1], 1) + elif t.dim() == 1: # (num_envs,) + t = t.view(-1, 1, 1) + elif t.dim() == 2: # (num_envs, 1) + t = t.unsqueeze(-1) + + if t.shape[1] == 1 and q0.shape[1] > 1: + t = t.repeat(1, q0.shape[1], 1) + + EPS = 1e-6 + + d = torch.sum(q0 * q1, dim=-1, keepdim=True) # (num_envs, bodys, 1) + + zero_mask = torch.isclose(t, torch.zeros_like(t), atol=EPS) + ones_mask = torch.isclose(t, torch.ones_like(t), atol=EPS) + dist_mask = (torch.abs(torch.abs(d) - 1.0) < EPS) + + out = torch.zeros_like(q0) + out[zero_mask.squeeze(-1)] = q0[zero_mask.squeeze(-1)] + out[ones_mask.squeeze(-1)] = q1[ones_mask.squeeze(-1)] + out[dist_mask.squeeze(-1)] = q0[dist_mask.squeeze(-1)] + + if shortestpath: + q1 = torch.where(d < 0, -q1, q1) + d = torch.abs(d) + + angle = torch.acos(torch.clamp(d, -1.0 + EPS, 1.0 - EPS)) + angle_mask = (torch.abs(angle) < EPS) + out[angle_mask.squeeze(-1)] = q0[angle_mask.squeeze(-1)] + + final_mask = ~(zero_mask | ones_mask | dist_mask | angle_mask) + final_mask = final_mask.squeeze(-1) + + if final_mask.any(): + sin_angle = torch.sin(angle) + isin = 1.0 / (sin_angle + EPS) + + t_expanded = t.expand_as(q0) + weight0 = torch.sin((1.0 - t_expanded) * angle) * isin + weight1 = torch.sin(t_expanded * angle) * isin + + result = weight0 * q0 + weight1 * q1 + result_norm = torch.norm(result, dim=-1, keepdim=True) + result = result / (result_norm + EPS) + + out[final_mask] = result[final_mask] + + return out \ No newline at end of file diff --git a/src/mjlab_husky/lerobot_numpy.py b/src/mjlab_husky/lerobot_numpy.py new file mode 100644 index 0000000..d43b87e --- /dev/null +++ b/src/mjlab_husky/lerobot_numpy.py @@ -0,0 +1,43 @@ +"""NumPy conversions for LeRobot v3 frames (batched torch / mjlab arrays).""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch + + +def import_lerobot_dataset_class(): + """Return ``LeRobotDataset`` or raise with install hints (optional ``lerobot`` extra).""" + try: + from lerobot.datasets.lerobot_dataset import LeRobotDataset # type: ignore + except Exception as e: # pragma: no cover + raise RuntimeError( + "未安装 `lerobot`（LeRobot v3 写数据需要）。\n" + "在项目目录执行：\n" + " uv pip install lerobot\n" + f"原始导入错误: {e}" + ) from e + return LeRobotDataset + + +def batched_vector_to_numpy1d(x: Any, env_i: int = 0) -> np.ndarray: + """Take env index `env_i`, flatten to 1D float32 (LeRobot parquet-friendly).""" + if x is None: + raise ValueError("x is None") + if isinstance(x, torch.Tensor): + t = x[env_i].detach().float().cpu().numpy() + elif hasattr(x, "detach"): + t = x[env_i].detach().float().cpu().numpy() + elif hasattr(x, "__getitem__") and hasattr(x, "cpu"): + t = x[env_i].float().cpu().numpy() + else: + arr = np.asarray(x, dtype=np.float32) + if arr.ndim >= 2: + arr = arr[env_i] + t = np.asarray(arr, dtype=np.float32) + out = np.asarray(t, dtype=np.float32).reshape(-1) + if out.size == 0: + raise RuntimeError("empty vector after flatten") + return out diff --git a/src/mjlab_husky/mujoco_gl.py b/src/mjlab_husky/mujoco_gl.py new file mode 100644 index 0000000..b682dc1 --- /dev/null +++ b/src/mjlab_husky/mujoco_gl.py @@ -0,0 +1,30 @@ +"""MuJoCo OpenGL backend: headless Linux needs EGL (or OSMesa), not GLFW.""" + +from __future__ import annotations + +import os +import sys + + +def ensure_mujoco_headless_gl() -> None: + """If no display server, set GL backends *before* ``import mujoco`` / ``import OpenGL``. + + On headless Linux, **PyOpenGL** would still match ``sys.platform == "linux"`` to the + **GLX** plugin, so ``PLATFORM.GL`` stays ``None`` and you get errors like + ``'NoneType' has no attribute 'glGetError'`` even when ``MUJOCO_GL=osmesa``. + + We set ``PYOPENGL_PLATFORM`` to match ``MUJOCO_GL`` (``osmesa`` or ``egl``). + """ + if sys.platform != "linux": + return + if os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY"): + return + + if not os.environ.get("MUJOCO_GL"): + os.environ.setdefault("MUJOCO_GL", "osmesa") + + gl = (os.environ.get("MUJOCO_GL") or "osmesa").lower() + if gl == "osmesa": + os.environ.setdefault("PYOPENGL_PLATFORM", "osmesa") + elif gl == "egl": + os.environ.setdefault("PYOPENGL_PLATFORM", "egl") diff --git a/src/mjlab_husky/rl/__init__.py b/src/mjlab_husky/rl/__init__.py new file mode 100644 index 0000000..a87bd24 --- /dev/null +++ b/src/mjlab_husky/rl/__init__.py @@ -0,0 +1,6 @@ +from mjlab.rl.config import RslRlBaseRunnerCfg as RslRlBaseRunnerCfg +from mjlab.rl.config import RslRlOnPolicyRunnerCfg as RslRlOnPolicyRunnerCfg +from mjlab.rl.config import RslRlPpoActorCriticCfg as RslRlPpoActorCriticCfg +from mjlab.rl.config import RslRlPpoAlgorithmCfg as RslRlPpoAlgorithmCfg +from mjlab_husky.rl.vecenv_wrapper import RslRlVecEnvWrapper as RslRlVecEnvWrapper +from mjlab_husky.rl.config import RslRlAMPOnPolicyRunnerCfg as RslRlAMPOnPolicyRunnerCfg \ No newline at end of file diff --git a/src/mjlab_husky/rl/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/rl/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..54767bc Binary files /dev/null and b/src/mjlab_husky/rl/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/mjlab_husky/rl/__pycache__/config.cpython-312.pyc b/src/mjlab_husky/rl/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000..d83a30a Binary files /dev/null and b/src/mjlab_husky/rl/__pycache__/config.cpython-312.pyc differ diff --git a/src/mjlab_husky/rl/__pycache__/vecenv_wrapper.cpython-312.pyc b/src/mjlab_husky/rl/__pycache__/vecenv_wrapper.cpython-312.pyc new file mode 100644 index 0000000..2bc380e Binary files /dev/null and b/src/mjlab_husky/rl/__pycache__/vecenv_wrapper.cpython-312.pyc differ diff --git a/src/mjlab_husky/rl/config.py b/src/mjlab_husky/rl/config.py new file mode 100644 index 0000000..006ad79 --- /dev/null +++ b/src/mjlab_husky/rl/config.py @@ -0,0 +1,124 @@ +"""RSL-RL configuration.""" + +from dataclasses import dataclass, field +from typing import Literal, Tuple + + +@dataclass +class RslRlPpoActorCriticCfg: + """Config for the PPO actor-critic networks.""" + + init_noise_std: float = 1.0 + """The initial noise standard deviation of the policy.""" + noise_std_type: Literal["scalar", "log"] = "scalar" + """The type of noise standard deviation for the policy. Default is scalar.""" + actor_obs_normalization: bool = False + """Whether to normalize the observation for the actor network. Default is False.""" + critic_obs_normalization: bool = False + """Whether to normalize the observation for the critic network. Default is False.""" + actor_hidden_dims: Tuple[int, ...] = (128, 128, 128) + """The hidden dimensions of the actor network.""" + critic_hidden_dims: Tuple[int, ...] = (128, 128, 128) + """The hidden dimensions of the critic network.""" + activation: str = "elu" + """The activation function to use in the actor and critic networks.""" + class_name: str = "ActorCritic" + """Ignore, required by RSL-RL.""" + + +@dataclass +class RslRlPpoAlgorithmCfg: + """Config for the PPO algorithm.""" + + num_learning_epochs: int = 5 + """The number of learning epochs per update.""" + num_mini_batches: int = 4 + """The number of mini-batches per update. + mini batch size = num_envs * num_steps / num_mini_batches + """ + learning_rate: float = 1e-3 + """The learning rate.""" + schedule: Literal["adaptive", "fixed"] = "adaptive" + """The learning rate schedule.""" + gamma: float = 0.99 + """The discount factor.""" + lam: float = 0.95 + """The lambda parameter for Generalized Advantage Estimation (GAE).""" + entropy_coef: float = 0.005 + """The coefficient for the entropy loss.""" + desired_kl: float = 0.01 + """The desired KL divergence between the new and old policies.""" + max_grad_norm: float = 1.0 + """The maximum gradient norm for the policy.""" + value_loss_coef: float = 1.0 + """The coefficient for the value loss.""" + use_clipped_value_loss: bool = True + """Whether to use clipped value loss.""" + clip_param: float = 0.2 + """The clipping parameter for the policy.""" + normalize_advantage_per_mini_batch: bool = False + """Whether to normalize the advantage per mini-batch. Default is False. If True, the + advantage is normalized over the mini-batches only. Otherwise, the advantage is + normalized over the entire collected trajectories. + """ + class_name: str = "PPO" + """Ignore, required by RSL-RL.""" + + +@dataclass +class RslRlBaseRunnerCfg: + seed: int = 42 + """The seed for the experiment. Default is 42.""" + num_steps_per_env: int = 24 + """The number of steps per environment update.""" + max_iterations: int = 300 + """The maximum number of iterations.""" + obs_groups: dict[str, tuple[str, ...]] = field( + default_factory=lambda: {"policy": ("policy",), "critic": ("critic",)}, + ) + save_interval: int = 50 + """The number of iterations between saves.""" + experiment_name: str = "exp1" + """The experiment name.""" + run_name: str = "" + """The run name. Default is empty string.""" + logger: Literal["wandb", "tensorboard"] = "wandb" + """The logger to use. Default is wandb.""" + wandb_project: str = "mjlab" + """The wandb project name.""" + wandb_tags: Tuple[str, ...] = () + """Tags for the wandb run. Default is empty tuple.""" + resume: bool = False + """Whether to resume the experiment. Default is False.""" + load_run: str = ".*" + """The run directory to load. Default is ".*" which means all runs. If regex + expression, the latest (alphabetical order) matching run will be loaded. + """ + load_checkpoint: str = "model_.*.pt" + """The checkpoint file to load. Default is "model_.*.pt" (all). If regex expression, + the latest (alphabetical order) matching file will be loaded. + """ + clip_actions: float | None = None + """The clipping range for action values. If None (default), no clipping is applied.""" + + +@dataclass +class RslRlOnPolicyRunnerCfg(RslRlBaseRunnerCfg): + class_name: str = "OnPolicyRunner" + """The runner class name. Default is OnPolicyRunner.""" + policy: RslRlPpoActorCriticCfg = field(default_factory=RslRlPpoActorCriticCfg) + """The policy configuration.""" + algorithm: RslRlPpoAlgorithmCfg = field(default_factory=RslRlPpoAlgorithmCfg) + """The algorithm configuration.""" + +@dataclass +class RslRlAMPOnPolicyRunnerCfg(RslRlOnPolicyRunnerCfg): + amp_num_obs: int = 23 + amp_num_frames: int = 5 + use_lerp: bool = False + amp_task_reward_lerp: float = 0.7 + amp_reward_coef: float = 5.0 + amp_motion_files: str = "dataset/skate_push" + amp_num_preload_transitions: int = 200000 + amp_discr_hidden_dims: Tuple[int, ...] = (256, 256) + min_normalized_std: Tuple[float, ...] = (0.05,) * 20 \ No newline at end of file diff --git a/src/mjlab_husky/rl/exporter_utils.py b/src/mjlab_husky/rl/exporter_utils.py new file mode 100644 index 0000000..0943f6d --- /dev/null +++ b/src/mjlab_husky/rl/exporter_utils.py @@ -0,0 +1,82 @@ +"""Shared utilities for ONNX policy export across RL tasks.""" + +import onnx +import torch + +from mjlab.entity import Entity +from mjlab.envs import ManagerBasedRlEnv +from mjlab.envs.mdp.actions import JointPositionAction + + +def list_to_csv_str(arr, *, decimals: int = 3, delimiter: str = ",") -> str: + """Convert list to CSV string with specified decimal precision.""" + fmt = f"{{:.{decimals}f}}" + return delimiter.join( + fmt.format(x) + if isinstance(x, (int, float)) + else str(x) # numbers → format, strings → as-is + for x in arr + ) + + +def get_base_metadata( + env: ManagerBasedRlEnv, run_path: str +) -> dict[str, list | str | float]: + """Get base metadata common to all RL policy exports. + + Args: + env: The RL environment. + run_path: W&B run path or other identifier. + + Returns: + Dictionary of metadata fields that are common across all tasks. + """ + robot: Entity = env.scene["robot"] + joint_action = env.action_manager.get_term("joint_pos") + assert isinstance(joint_action, JointPositionAction) + # Build mapping from joint name to actuator ID for natural joint order. + # Each spec actuator controls exactly one joint (via its target field). + joint_name_to_ctrl_id = {} + for actuator in robot.spec.actuators: + joint_name = actuator.target.split("/")[-1] + joint_name_to_ctrl_id[joint_name] = actuator.id + # Get actuator IDs in natural joint order (same order as robot.joint_names). + ctrl_ids_natural = [ + joint_name_to_ctrl_id[jname] + for jname in robot.joint_names # global joint order + if jname in joint_name_to_ctrl_id # skip non-actuated joints + ] + joint_stiffness = env.sim.mj_model.actuator_gainprm[ctrl_ids_natural, 0] + joint_damping = -env.sim.mj_model.actuator_biasprm[ctrl_ids_natural, 2] + return { + "run_path": run_path, + "joint_names": list(robot.joint_names), + "joint_stiffness": joint_stiffness.tolist(), + "joint_damping": joint_damping.tolist(), + "default_joint_pos": robot.data.default_joint_pos[0].cpu().tolist(), + "command_names": list(env.command_manager.active_terms), + "observation_names": env.observation_manager.active_terms["policy"], + "action_scale": joint_action._scale[0].cpu().tolist() + if isinstance(joint_action._scale, torch.Tensor) + else joint_action._scale, + } + + +def attach_metadata_to_onnx( + onnx_path: str, metadata: dict[str, list | str | float] +) -> None: + """Attach metadata to an ONNX model file. + + Args: + onnx_path: Path to the ONNX model file. + metadata: Dictionary of metadata key-value pairs to attach. + """ + model = onnx.load(onnx_path) + + for k, v in metadata.items(): + entry = onnx.StringStringEntryProto() + entry.key = k + entry.value = list_to_csv_str(v) if isinstance(v, list) else str(v) + model.metadata_props.append(entry) + + onnx.save(model, onnx_path) diff --git a/src/mjlab_husky/rl/vecenv_wrapper.py b/src/mjlab_husky/rl/vecenv_wrapper.py new file mode 100644 index 0000000..98b164f --- /dev/null +++ b/src/mjlab_husky/rl/vecenv_wrapper.py @@ -0,0 +1,127 @@ +import torch +from rsl_rl.env import VecEnv +from tensordict import TensorDict + +from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg +from mjlab.utils.spaces import Space + + +class RslRlVecEnvWrapper(VecEnv): + def __init__( + self, + env: ManagerBasedRlEnv, + clip_actions: float | None = None, + ): + self.env = env + self.clip_actions = clip_actions + + self.num_envs = self.unwrapped.num_envs + self.device = torch.device(self.unwrapped.device) + self.max_episode_length = self.unwrapped.max_episode_length + self.max_episode_length_s = self.unwrapped.max_episode_length_s + self.num_actions = self.unwrapped.action_manager.total_action_dim + self._modify_action_space() + + # Reset at the start since rsl_rl does not call reset. + self.env.reset() + + @property + def cfg(self) -> ManagerBasedRlEnvCfg: + return self.unwrapped.cfg + + @property + def render_mode(self) -> str | None: + return self.env.render_mode + + @property + def observation_space(self) -> Space: + return self.env.observation_space + + @property + def action_space(self) -> Space: + return self.env.action_space + + @classmethod + def class_name(cls) -> str: + return cls.__name__ + + @property + def unwrapped(self) -> ManagerBasedRlEnv: + return self.env + + # Properties. + + @property + def episode_length_buf(self) -> torch.Tensor: + return self.unwrapped.episode_length_buf + + @property + def reset_env_ids(self) -> torch.Tensor | None: + return self.unwrapped.reset_env_ids + + @reset_env_ids.setter + def reset_env_ids(self, value: torch.Tensor | None) -> None: + self.unwrapped.reset_env_ids = value + + @property + def contact_phase(self) -> torch.Tensor | None: + return self.unwrapped.contact_phase + + @contact_phase.setter + def contact_phase(self, value: torch.Tensor | None) -> None: + self.unwrapped.contact_phase = value + + @episode_length_buf.setter + def episode_length_buf(self, value: torch.Tensor) -> None: # type: ignore + self.unwrapped.episode_length_buf = value + + def seed(self, seed: int = -1) -> int: + return self.unwrapped.seed(seed) + + def get_observations(self) -> TensorDict: + obs_dict = self.unwrapped.observation_manager.compute() + return TensorDict(obs_dict, batch_size=[self.num_envs]) + + def reset(self) -> tuple[TensorDict, dict]: + obs_dict, extras = self.env.reset() + return TensorDict(obs_dict, batch_size=[self.num_envs]), extras + + def step( + self, actions: torch.Tensor + ) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]: + if self.clip_actions is not None: + actions = torch.clamp(actions, -self.clip_actions, self.clip_actions) + obs_dict, rew, terminated, truncated, extras = self.env.step(actions) + term_or_trunc = terminated | truncated + assert isinstance(rew, torch.Tensor) + assert isinstance(term_or_trunc, torch.Tensor) + dones = term_or_trunc.to(dtype=torch.long) + if not self.cfg.is_finite_horizon: + extras["time_outs"] = truncated + return ( + TensorDict(obs_dict, batch_size=[self.num_envs]), + rew, + dones, + extras, + ) + + def close(self) -> None: + return self.env.close() + + # Private methods. + + def _modify_action_space(self) -> None: + if self.clip_actions is None: + return + + from mjlab.utils.spaces import Box, batch_space + + self.unwrapped.single_action_space = Box( + shape=(self.num_actions,), low=-self.clip_actions, high=self.clip_actions + ) + self.unwrapped.action_space = batch_space( + self.unwrapped.single_action_space, self.num_envs + ) + + def get_amp_observations(self) -> TensorDict: + return self.unwrapped.get_amp_observations() \ No newline at end of file diff --git a/src/mjlab_husky/scripts/__pycache__/play.cpython-312.pyc b/src/mjlab_husky/scripts/__pycache__/play.cpython-312.pyc new file mode 100644 index 0000000..acc2411 Binary files /dev/null and b/src/mjlab_husky/scripts/__pycache__/play.cpython-312.pyc differ diff --git a/src/mjlab_husky/scripts/__pycache__/play.cpython-313.pyc b/src/mjlab_husky/scripts/__pycache__/play.cpython-313.pyc new file mode 100644 index 0000000..aca5d8b Binary files /dev/null and b/src/mjlab_husky/scripts/__pycache__/play.cpython-313.pyc differ diff --git a/src/mjlab_husky/scripts/export_lerobot_qpos.py b/src/mjlab_husky/scripts/export_lerobot_qpos.py new file mode 100644 index 0000000..21ed84c --- /dev/null +++ b/src/mjlab_husky/scripts/export_lerobot_qpos.py @@ -0,0 +1,140 @@ +"""Export mjlab_husky trajectories to LeRobotDataset v3 (parquet). + +Each frame stores ``observation.state`` (qpos) and ``action`` (policy output for ``env.step``). +""" + +from __future__ import annotations + +import shutil +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Literal + +import numpy as np +import torch +import tyro + +from mjlab_husky.lerobot_numpy import batched_vector_to_numpy1d, import_lerobot_dataset_class +from mjlab_husky.mujoco_gl import ensure_mujoco_headless_gl + +ensure_mujoco_headless_gl() + +from mjlab.utils.torch import configure_torch_backends +from mjlab_husky.envs import G1SkaterManagerBasedRlEnv +from mjlab_husky.rl import RslRlVecEnvWrapper +from mjlab_husky.tasks.registry import load_env_cfg, load_rl_cfg, load_runner_cls + + +@dataclass(frozen=True) +class ExportConfig: + task_id: str + checkpoint_file: str + out_dir: str = "lerobot_data" + dataset_name: str = "mjlab_husky_skater_qpos" + agent: Literal["trained", "zero", "random"] = "trained" + device: str | None = None + num_envs: int = 1 + episodes: int = 1 + steps_per_episode: int = 1000 + fps: int | None = None + task_text: str = "Skater" + overwrite: bool = False + + +def main(cfg: ExportConfig) -> None: + configure_torch_backends() + LeRobotDataset = import_lerobot_dataset_class() + + device = cfg.device or ("cuda:0" if torch.cuda.is_available() else "cpu") + env_cfg = load_env_cfg(cfg.task_id, play=True) + agent_cfg = load_rl_cfg(cfg.task_id) + + env_cfg.scene.num_envs = int(cfg.num_envs) + env = G1SkaterManagerBasedRlEnv(cfg=env_cfg, device=device, render_mode=None) + env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions) + + action_shape: tuple[int, ...] = env.unwrapped.action_space.shape # type: ignore + if cfg.agent == "zero": + + class PolicyZero: + def __call__(self, obs) -> torch.Tensor: + del obs + return torch.zeros(action_shape, device=env.unwrapped.device) + + policy = PolicyZero() + elif cfg.agent == "random": + + class PolicyRandom: + def __call__(self, obs) -> torch.Tensor: + del obs + return 2 * torch.rand(action_shape, device=env.unwrapped.device) - 1 + + policy = PolicyRandom() + else: + runner_cls = load_runner_cls(cfg.task_id) + runner = runner_cls(env, asdict(agent_cfg), device=device) + runner.load(str(Path(cfg.checkpoint_file)), map_location=device) + policy = runner.get_inference_policy(device=device) + + reset_out = env.reset() + obs = reset_out[0] if isinstance(reset_out, tuple) and len(reset_out) == 2 else reset_out + sim = getattr(getattr(env.unwrapped, "sim", None), "data", None) + if sim is None: + raise RuntimeError("env.unwrapped.sim.data not found; cannot read qpos") + qpos = getattr(sim, "qpos", None) + if qpos is None: + raise RuntimeError("sim.data.qpos not found") + q0 = batched_vector_to_numpy1d(qpos, env_i=0) + nq = int(q0.shape[0]) + na = int(np.prod(np.asarray(env.unwrapped.action_space.shape, dtype=np.int64))) + + step_dt = float(getattr(env.unwrapped, "step_dt", 0.02)) + fps = int(cfg.fps) if cfg.fps is not None else max(1, int(round(1.0 / step_dt))) + + out_root = (Path.cwd() / cfg.out_dir).resolve() + ds_root = out_root / cfg.dataset_name + if ds_root.exists(): + if not cfg.overwrite: + raise FileExistsError(f"Dataset folder exists: {ds_root} (use --overwrite)") + shutil.rmtree(ds_root) + ds_root.parent.mkdir(parents=True, exist_ok=True) + + features = { + "observation.state": {"dtype": "float32", "shape": (nq,), "names": None}, + "action": {"dtype": "float32", "shape": (na,), "names": None}, + } + dataset = LeRobotDataset.create( + repo_id=cfg.dataset_name, + fps=fps, + features=features, + root=str(ds_root), + robot_type="mjlab_husky", + use_videos=False, + ) + + for _ep in range(int(cfg.episodes)): + reset_out = env.reset() + obs = reset_out[0] if isinstance(reset_out, tuple) and len(reset_out) == 2 else reset_out + + for _ in range(int(cfg.steps_per_episode)): + sim = env.unwrapped.sim.data # type: ignore[attr-defined] + q = batched_vector_to_numpy1d(sim.qpos, env_i=0) + action = policy(obs) + a = batched_vector_to_numpy1d(action, env_i=0) + dataset.add_frame({"task": cfg.task_text, "observation.state": q, "action": a}) + step_out = env.step(action) + obs = step_out[0] + + dataset.save_episode() + + dataset.finalize() + env.close() + + print(f"[OK] LeRobot v3 dataset (state+action) written to: {ds_root}") + + +if __name__ == "__main__": + import mjlab.tasks # noqa: F401 + + cfg = tyro.cli(ExportConfig) + main(cfg) diff --git a/src/mjlab_husky/scripts/play.py b/src/mjlab_husky/scripts/play.py new file mode 100644 index 0000000..cad1ae7 --- /dev/null +++ b/src/mjlab_husky/scripts/play.py @@ -0,0 +1,335 @@ +"""Script to play RL agent with RSL-RL.""" + +from __future__ import annotations + +import os +import sys +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Literal + +from mjlab_husky.mujoco_gl import ensure_mujoco_headless_gl + +ensure_mujoco_headless_gl() + +import torch +import tyro + +from mjlab.envs import ManagerBasedRlEnvCfg, ManagerBasedRlEnv +from mjlab_husky.envs import G1SkaterManagerBasedRlEnvCfg, G1SkaterManagerBasedRlEnv +from mjlab_husky.rl import RslRlVecEnvWrapper +from mjlab_husky.tasks.registry import list_tasks, load_env_cfg, load_rl_cfg, load_runner_cls +from mjlab.utils.os import get_wandb_checkpoint_path +from mjlab.utils.torch import configure_torch_backends +from mjlab.utils.wrappers import VideoRecorder +from mjlab.viewer import NativeMujocoViewer, ViserPlayViewer +from mjlab_husky.viewer.rerun_native_play_viewer import RerunNativePlayViewer +from mjlab_husky.viewer.rerun_play_viewer import RerunPlayViewer, RerunPlayViewerCfg +from mjlab_husky.viewer.rerun_viser_play_viewer import RerunViserPlayViewer + + +@dataclass(frozen=True) +class PlayConfig: + agent: Literal["zero", "random", "trained"] = "trained" + registry_name: str | None = None + wandb_run_path: str | None = None + checkpoint_file: str | None = None + motion_file: str | None = None + num_envs: int | None = None + device: str | None = None + video: bool = False + video_length: int = 1000 + video_height: int = 1080 + video_width: int = 1920 + camera: int | str | None = None + viewer: Literal["auto", "native", "viser", "rerun", "rerun_native", "rerun_viser"] = "auto" + rerun_web_port: int = 8080 + rerun_grpc_port: int | None = None + # Viser 独立端口（与 rerun_web_port 分离，便于 RoboHub 双 iframe / 避免默认都与 8080 冲突） + viser_port: int = 8081 + # 浏览器直连云主机时填公网/LAN IP（与 --rerun-web-port 对应）；否则用 SSH 转发保持默认 127.0.0.1。 + rerun_connect_host: str | None = None + # If False, open the printed `?url=...` link yourself. Use `--no-rerun-open-browser`. + rerun_open_browser: bool = True + # Rerun-only offscreen size (when not --video). 1080p every step will stall gRPC / viewer. + rerun_viewer_height: int = 540 + rerun_viewer_width: int = 960 + # Log camera every N env steps; qpos every M steps (cuts bandwidth / memory). + rerun_camera_log_stride: int = 2 + rerun_qpos_log_stride: int = 4 + # Shrink image before rr.log (longest side, pixels). + rerun_camera_max_side: int = 640 + # Prefer fresh frames in viewer when the client falls behind. + rerun_grpc_newest_first: bool = True + rerun_log_qpos_joint_scalars: bool = True + rerun_log_qpos_text: bool = False + rerun_qpos_text_decimals: int = 3 + rerun_log_qpos_tensor: bool = False + rerun_qpos_env_idx: int = 0 + + # --- LeRobot v3: record qpos while playing Rerun --- + lerobot_record: bool = False + lerobot_out_dir: str = "lerobot_data" + lerobot_dataset_name: str = "mjlab_husky_live" + lerobot_steps_per_episode: int = 1000 + lerobot_task: str = "Skater" + lerobot_overwrite: bool = False + lerobot_record_stride: int = 1 + # 每次仿真 done->reset 就存一个 LeRobot episode（与「录满 N 帧再存」二选一或叠加使用） + lerobot_save_on_env_reset: bool = False + + # Internal flag used by demo script. + _demo_mode: tyro.conf.Suppress[bool] = False + + +def _resolve_play_viewer(cfg: PlayConfig) -> str: + if cfg.viewer == "auto": + has_display = bool(os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY")) + return "native" if has_display else "rerun" + return cfg.viewer + + +def run_play(task_id: str, cfg: PlayConfig): + configure_torch_backends() + + device = cfg.device or ("cuda:0" if torch.cuda.is_available() else "cpu") + + env_cfg = load_env_cfg(task_id, play=True) + agent_cfg = load_rl_cfg(task_id) + + DUMMY_MODE = cfg.agent in {"zero", "random"} + TRAINED_MODE = not DUMMY_MODE + + log_dir: Path | None = None + resume_path: Path | None = None + if TRAINED_MODE: + log_root_path = (Path("logs") / "rsl_rl" / agent_cfg.experiment_name).resolve() + if cfg.checkpoint_file is not None: + resume_path = Path(cfg.checkpoint_file) + if not resume_path.exists(): + raise FileNotFoundError(f"Checkpoint file not found: {resume_path}") + print(f"[INFO]: Loading checkpoint: {resume_path.name}") + else: + if cfg.wandb_run_path is None: + raise ValueError( + "`wandb_run_path` is required when `checkpoint_file` is not provided." + ) + resume_path, was_cached = get_wandb_checkpoint_path( + log_root_path, Path(cfg.wandb_run_path) + ) + # Extract run_id and checkpoint name from path for display. + run_id = resume_path.parent.name + checkpoint_name = resume_path.name + cached_str = "cached" if was_cached else "downloaded" + print( + f"[INFO]: Loading checkpoint: {checkpoint_name} (run: {run_id}, {cached_str})" + ) + log_dir = resume_path.parent + + if cfg.num_envs is not None: + env_cfg.scene.num_envs = cfg.num_envs + + resolved_viewer = _resolve_play_viewer(cfg) + if resolved_viewer in ("rerun", "rerun_native", "rerun_viser") and not cfg.video: + env_cfg.viewer.height = cfg.rerun_viewer_height + env_cfg.viewer.width = cfg.rerun_viewer_width + else: + env_cfg.viewer.height = cfg.video_height + env_cfg.viewer.width = cfg.video_width + + # Rerun needs rgb_array so we can log the same offscreen camera as mjlab (tracking robot). + render_mode = ( + "rgb_array" + if ( + (TRAINED_MODE and cfg.video) + or resolved_viewer in ("rerun", "rerun_native", "rerun_viser") + ) + else None + ) + if cfg.video and DUMMY_MODE: + print( + "[WARN] Video recording with dummy agents is disabled (no checkpoint/log_dir)." + ) + env = G1SkaterManagerBasedRlEnv(cfg=env_cfg, device=device, render_mode=render_mode) + + if TRAINED_MODE and cfg.video: + print("[INFO] Recording videos during play") + assert log_dir is not None # log_dir is set in TRAINED_MODE block + env = VideoRecorder( + env, + video_folder=log_dir / "videos" / "play", + step_trigger=lambda step: step == 0, + video_length=cfg.video_length, + disable_logger=True, + ) + + env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions) + if DUMMY_MODE: + action_shape: tuple[int, ...] = env.unwrapped.action_space.shape # type: ignore + if cfg.agent == "zero": + + class PolicyZero: + def __call__(self, obs) -> torch.Tensor: + del obs + return torch.zeros(action_shape, device=env.unwrapped.device) + + policy = PolicyZero() + else: + + class PolicyRandom: + def __call__(self, obs) -> torch.Tensor: + del obs + return 2 * torch.rand(action_shape, device=env.unwrapped.device) - 1 + + policy = PolicyRandom() + else: + runner_cls = load_runner_cls(task_id) + runner = runner_cls(env, asdict(agent_cfg), device=device) + runner.load(str(resume_path), map_location=device) + policy = runner.get_inference_policy(device=device) + + if cfg.lerobot_record and resolved_viewer in ("rerun_native", "rerun_viser"): + raise RuntimeError( + "`--lerobot-record` 当前仅支持 `--viewer rerun`（独立 Rerun 循环）。" + "请去掉 `--lerobot-record` 或改用 `--viewer rerun`。" + ) + + if resolved_viewer == "rerun_viser": + if cfg.viser_port == cfg.rerun_web_port: + raise ValueError( + "`--viser-port` 与 `--rerun-web-port` 不能相同（Viser 与 Rerun Web 需各占一个端口）。" + ) + if cfg.rerun_grpc_port is not None and cfg.viser_port == cfg.rerun_grpc_port: + raise ValueError( + "`--viser-port` 与 `--rerun-grpc-port` 不能相同。" + ) + + if resolved_viewer == "native": + NativeMujocoViewer(env, policy).run() + elif resolved_viewer == "viser": + ViserPlayViewer(env, policy).run() + elif resolved_viewer == "rerun": + RerunPlayViewer( + env, + policy, + cfg=RerunPlayViewerCfg( + web_port=cfg.rerun_web_port, + open_browser=cfg.rerun_open_browser, + grpc_port=cfg.rerun_grpc_port, + connect_host=cfg.rerun_connect_host, + log_camera=True, + camera_log_stride=max(1, cfg.rerun_camera_log_stride), + qpos_log_stride=max(1, cfg.rerun_qpos_log_stride), + camera_max_side=max(64, cfg.rerun_camera_max_side), + grpc_newest_first=cfg.rerun_grpc_newest_first, + log_qpos_joint_scalars=cfg.rerun_log_qpos_joint_scalars, + log_qpos_text=cfg.rerun_log_qpos_text, + qpos_text_decimals=max(0, min(8, cfg.rerun_qpos_text_decimals)), + log_qpos_tensor=cfg.rerun_log_qpos_tensor, + qpos_env_idx=max(0, cfg.rerun_qpos_env_idx), + lerobot_record=cfg.lerobot_record, + lerobot_out_dir=cfg.lerobot_out_dir, + lerobot_dataset_name=cfg.lerobot_dataset_name, + lerobot_steps_per_episode=max(1, cfg.lerobot_steps_per_episode), + lerobot_task=cfg.lerobot_task, + lerobot_overwrite=cfg.lerobot_overwrite, + lerobot_record_stride=max(1, cfg.lerobot_record_stride), + lerobot_save_on_env_reset=cfg.lerobot_save_on_env_reset, + ), + ).run() + elif resolved_viewer == "rerun_native": + RerunNativePlayViewer( + env, + policy, + rerun_cfg=RerunPlayViewerCfg( + web_port=cfg.rerun_web_port, + open_browser=cfg.rerun_open_browser, + grpc_port=cfg.rerun_grpc_port, + connect_host=cfg.rerun_connect_host, + log_camera=True, + camera_log_stride=max(1, cfg.rerun_camera_log_stride), + qpos_log_stride=max(1, cfg.rerun_qpos_log_stride), + camera_max_side=max(64, cfg.rerun_camera_max_side), + grpc_newest_first=cfg.rerun_grpc_newest_first, + log_qpos_joint_scalars=cfg.rerun_log_qpos_joint_scalars, + log_qpos_text=cfg.rerun_log_qpos_text, + qpos_text_decimals=max(0, min(8, cfg.rerun_qpos_text_decimals)), + log_qpos_tensor=cfg.rerun_log_qpos_tensor, + qpos_env_idx=max(0, cfg.rerun_qpos_env_idx), + lerobot_record=False, + lerobot_out_dir=cfg.lerobot_out_dir, + lerobot_dataset_name=cfg.lerobot_dataset_name, + lerobot_steps_per_episode=max(1, cfg.lerobot_steps_per_episode), + lerobot_task=cfg.lerobot_task, + lerobot_overwrite=cfg.lerobot_overwrite, + lerobot_record_stride=max(1, cfg.lerobot_record_stride), + lerobot_save_on_env_reset=cfg.lerobot_save_on_env_reset, + ), + ).run() + elif resolved_viewer == "rerun_viser": + RerunViserPlayViewer( + env, + policy, + rerun_cfg=RerunPlayViewerCfg( + web_port=cfg.rerun_web_port, + open_browser=cfg.rerun_open_browser, + grpc_port=cfg.rerun_grpc_port, + connect_host=cfg.rerun_connect_host, + log_camera=True, + camera_log_stride=max(1, cfg.rerun_camera_log_stride), + qpos_log_stride=max(1, cfg.rerun_qpos_log_stride), + camera_max_side=max(64, cfg.rerun_camera_max_side), + grpc_newest_first=cfg.rerun_grpc_newest_first, + log_qpos_joint_scalars=cfg.rerun_log_qpos_joint_scalars, + log_qpos_text=cfg.rerun_log_qpos_text, + qpos_text_decimals=max(0, min(8, cfg.rerun_qpos_text_decimals)), + log_qpos_tensor=cfg.rerun_log_qpos_tensor, + qpos_env_idx=max(0, cfg.rerun_qpos_env_idx), + lerobot_record=False, + lerobot_out_dir=cfg.lerobot_out_dir, + lerobot_dataset_name=cfg.lerobot_dataset_name, + lerobot_steps_per_episode=max(1, cfg.lerobot_steps_per_episode), + lerobot_task=cfg.lerobot_task, + lerobot_overwrite=cfg.lerobot_overwrite, + lerobot_record_stride=max(1, cfg.lerobot_record_stride), + lerobot_save_on_env_reset=cfg.lerobot_save_on_env_reset, + ), + viser_port=cfg.viser_port, + ).run() + else: + raise RuntimeError(f"Unsupported viewer backend: {resolved_viewer}") + + env.close() + + +def main(): + # Parse first argument to choose the task. + # Import tasks to populate the registry. + import mjlab.tasks # noqa: F401 + + all_tasks = list_tasks() + chosen_task, remaining_args = tyro.cli( + tyro.extras.literal_type_from_choices(all_tasks), + add_help=False, + return_unknown_args=True, + ) + + # Parse the rest of the arguments + allow overriding env_cfg and agent_cfg. + agent_cfg = load_rl_cfg(chosen_task) + + args = tyro.cli( + PlayConfig, + args=remaining_args, + default=PlayConfig(), + prog=sys.argv[0] + f" {chosen_task}", + config=( + tyro.conf.AvoidSubcommands, + ), + ) + del remaining_args, agent_cfg + + run_play(chosen_task, args) + + +if __name__ == "__main__": + main() diff --git a/src/mjlab_husky/scripts/play_lerobot_rerun.py b/src/mjlab_husky/scripts/play_lerobot_rerun.py new file mode 100644 index 0000000..916b197 --- /dev/null +++ b/src/mjlab_husky/scripts/play_lerobot_rerun.py @@ -0,0 +1,100 @@ +"""一条命令：策略推理 + LeRobot v3 帧写入 + Rerun 实时播放。""" + +from __future__ import annotations + +from dataclasses import dataclass + +import tyro + +from mjlab_husky.mujoco_gl import ensure_mujoco_headless_gl + +ensure_mujoco_headless_gl() + + +@dataclass(frozen=True) +class PlayLeRobotRerunConfig: + task_id: str = "Mjlab-Skater-Flat-Unitree-G1" + checkpoint_file: str = "ckpts/test.pt" + device: str | None = None + num_envs: int | None = None + + lerobot_out_dir: str = "lerobot_data" + lerobot_dataset_name: str = "mjlab_husky_live" + lerobot_overwrite: bool = False + lerobot_steps_per_episode: int = 1000 + lerobot_task: str = "Skater" + lerobot_record_stride: int = 1 + lerobot_save_on_env_reset: bool = False + + rerun_web_port: int = 8080 + rerun_grpc_port: int | None = None + rerun_connect_host: str | None = None + rerun_open_browser: bool = False + rerun_viewer_width: int = 960 + rerun_viewer_height: int = 540 + rerun_camera_log_stride: int = 2 + rerun_qpos_log_stride: int = 4 + rerun_camera_max_side: int = 640 + rerun_grpc_newest_first: bool = True + rerun_log_qpos_joint_scalars: bool = False + rerun_log_qpos_text: bool = True + rerun_qpos_text_decimals: int = 3 + rerun_log_qpos_tensor: bool = False + rerun_qpos_env_idx: int = 0 + + +def _ensure_lerobot_installed() -> None: + try: + import lerobot # noqa: F401 + except ImportError as e: + raise SystemExit( + "未安装 lerobot，无法写入 LeRobot v3。\n在本仓库目录执行： uv pip install lerobot\n" + ) from e + + +def main() -> None: + c = tyro.cli(PlayLeRobotRerunConfig, prog="play-lerobot-rerun") + _ensure_lerobot_installed() + + import mjlab.tasks # noqa: F401 + + from mjlab_husky.scripts.play import PlayConfig, run_play + + print( + "[play-lerobot-rerun] 推理 + LeRobot v3 + Rerun；数据目录：" + f"{c.lerobot_out_dir}/{c.lerobot_dataset_name}" + ) + play_cfg = PlayConfig( + viewer="rerun", + lerobot_record=True, + checkpoint_file=c.checkpoint_file, + device=c.device, + num_envs=c.num_envs, + rerun_web_port=c.rerun_web_port, + rerun_grpc_port=c.rerun_grpc_port, + rerun_connect_host=c.rerun_connect_host, + rerun_open_browser=c.rerun_open_browser, + rerun_viewer_width=c.rerun_viewer_width, + rerun_viewer_height=c.rerun_viewer_height, + rerun_camera_log_stride=c.rerun_camera_log_stride, + rerun_qpos_log_stride=c.rerun_qpos_log_stride, + rerun_camera_max_side=c.rerun_camera_max_side, + rerun_grpc_newest_first=c.rerun_grpc_newest_first, + rerun_log_qpos_joint_scalars=c.rerun_log_qpos_joint_scalars, + rerun_log_qpos_text=c.rerun_log_qpos_text, + rerun_qpos_text_decimals=c.rerun_qpos_text_decimals, + rerun_log_qpos_tensor=c.rerun_log_qpos_tensor, + rerun_qpos_env_idx=c.rerun_qpos_env_idx, + lerobot_out_dir=c.lerobot_out_dir, + lerobot_dataset_name=c.lerobot_dataset_name, + lerobot_overwrite=c.lerobot_overwrite, + lerobot_steps_per_episode=c.lerobot_steps_per_episode, + lerobot_task=c.lerobot_task, + lerobot_record_stride=c.lerobot_record_stride, + lerobot_save_on_env_reset=c.lerobot_save_on_env_reset, + ) + run_play(c.task_id, play_cfg) + + +if __name__ == "__main__": + main() diff --git a/src/mjlab_husky/scripts/rerun_lerobot_dataset.py b/src/mjlab_husky/scripts/rerun_lerobot_dataset.py new file mode 100644 index 0000000..0d77551 --- /dev/null +++ b/src/mjlab_husky/scripts/rerun_lerobot_dataset.py @@ -0,0 +1,179 @@ +"""Stream a local LeRobot v3.0 dataset to the Rerun viewer (no MuJoCo).""" + +from __future__ import annotations + +import os +import re +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from urllib.parse import quote + +import numpy as np +import torch +import tyro + +import rerun as rr + + +def _viewer_url_with_connect(*, host: str, web_port: int, connect_uri: str) -> str: + return f"http://{host}:{web_port}/?url={quote(connect_uri, safe='')}" + + +def _grpc_listen_port(connect_uri: str) -> int | None: + m = re.search(r"127\.0\.0\.1:(\d+)", connect_uri) + return int(m.group(1)) if m else None + + +def _grpc_port_from_uri(connect_uri: str) -> int | None: + m = re.search(r":(\d+)/", connect_uri) + return int(m.group(1)) if m else None + + +def _rewrite_loopback_in_rerun_uri(uri: str, host: str) -> str: + return re.sub(r"127\.0\.0\.1", host.strip(), uri, count=1) + + +def _tensorish_to_numpy1d(v: Any) -> np.ndarray | None: + if v is None: + return None + try: + if isinstance(v, torch.Tensor): + t = v.detach().float().cpu().numpy().reshape(-1) + else: + t = np.asarray(v, dtype=np.float64).reshape(-1) + return np.asarray(t, dtype=np.float64).reshape(-1) + except Exception: + return None + + +def _log_vector_scalars(prefix: str, vec: np.ndarray) -> None: + for i, val in enumerate(vec): + rr.log(f"{prefix}/j_{i:03d}", rr.Scalars([float(val)])) + + +def _import_lerobot_dataset(): + try: + from lerobot.datasets.lerobot_dataset import LeRobotDataset # type: ignore + except Exception as e: # pragma: no cover + raise RuntimeError( + "需要安装 lerobot 才能读取 v3 数据集。\n uv pip install lerobot\n" f"原始错误: {e}" + ) from e + return LeRobotDataset + + +@dataclass(frozen=True) +class RerunLeRobotCfg: + root: Path + repo_id: str | None = None + web_port: int = 8080 + grpc_port: int | None = None + connect_host: str | None = None + open_browser: bool = True + grpc_newest_first: bool = True + download_videos: bool = False + frame_stride: int = 1 + max_frames: int | None = None + playback_fps: float | None = None + log_state_scalars: bool = True + log_action_scalars: bool = True + log_meta_scalars: bool = True + + +def run_rerun_lerobot(cfg: RerunLeRobotCfg) -> None: + root = cfg.root.expanduser().resolve() + if not (root / "meta" / "info.json").is_file(): + raise FileNotFoundError(f"不是有效的 LeRobot 数据集目录（缺少 meta/info.json）: {root}") + + repo_id = cfg.repo_id if cfg.repo_id is not None else root.name + LeRobotDataset = _import_lerobot_dataset() + ds = LeRobotDataset(repo_id=repo_id, root=root, download_videos=cfg.download_videos) + + rr.init("lerobot_dataset_rerun", spawn=False) + sg_kw: dict[str, Any] = {"newest_first": cfg.grpc_newest_first} + if cfg.grpc_port is not None: + sg_kw["grpc_port"] = cfg.grpc_port + server_uri = rr.serve_grpc(**sg_kw) + + connect_uri = server_uri + viewer_http_host = "127.0.0.1" + if cfg.connect_host: + connect_uri = _rewrite_loopback_in_rerun_uri(server_uri, cfg.connect_host) + viewer_http_host = cfg.connect_host.strip() + + local_viewer = _viewer_url_with_connect( + host=viewer_http_host, web_port=cfg.web_port, connect_uri=connect_uri + ) + grpc_port = _grpc_listen_port(server_uri) or _grpc_port_from_uri(connect_uri) + print("[Rerun] Open this URL:", local_viewer) + if grpc_port is not None: + print( + "[Rerun] SSH:", + f"ssh -N -L {cfg.web_port}:127.0.0.1:{cfg.web_port} " + f"-L {grpc_port}:127.0.0.1:{grpc_port} user@remote", + ) + + rr.serve_web_viewer( + web_port=cfg.web_port, + open_browser=cfg.open_browser, + connect_to=connect_uri, + ) + + n = len(ds) + limit = n if cfg.max_frames is None else min(n, int(cfg.max_frames)) + stride = max(1, int(cfg.frame_stride)) + sleep_s = (1.0 / float(cfg.playback_fps)) if cfg.playback_fps and cfg.playback_fps > 0 else 0.0 + + print(f"[LeRobot→Rerun] frames={n}, replaying up to {limit}, stride={stride}, root={root}") + + out_i = 0 + for idx in range(0, limit, stride): + rr.set_time("frame", sequence=out_i) + item = ds[idx] + + if cfg.log_meta_scalars: + for key in ("episode_index", "frame_index", "index", "timestamp"): + if key not in item: + continue + v = item[key] + try: + if isinstance(v, torch.Tensor): + val = float(v.detach().cpu().reshape(-1)[0].item()) + else: + val = float(np.asarray(v).reshape(-1)[0]) + rr.log(f"lerobot/{key}", rr.Scalars([val])) + except Exception: + pass + if "task" in item: + try: + t = item["task"] + if not isinstance(t, str): + t = str(t) + rr.log("lerobot/task", rr.TextLog(t)) + except Exception: + pass + + st = _tensorish_to_numpy1d(item.get("observation.state")) + if st is not None and cfg.log_state_scalars: + _log_vector_scalars("lerobot/observation_state", st) + + ac = _tensorish_to_numpy1d(item.get("action")) + if ac is not None and cfg.log_action_scalars: + _log_vector_scalars("lerobot/action", ac) + + out_i += 1 + if sleep_s > 0: + time.sleep(sleep_s) + + print("[LeRobot→Rerun] done.") + + +def main() -> None: + tyro.extras.set_accent_color("bright_yellow") + cfg = tyro.cli(RerunLeRobotCfg) + run_rerun_lerobot(cfg) + + +if __name__ == "__main__": + main() diff --git a/src/mjlab_husky/scripts/train.py b/src/mjlab_husky/scripts/train.py new file mode 100644 index 0000000..3b5dc3c --- /dev/null +++ b/src/mjlab_husky/scripts/train.py @@ -0,0 +1,223 @@ +"""Script to train RL agent with RSL-RL.""" + +import logging +import os +import sys +from dataclasses import asdict, dataclass, field +from datetime import datetime +from pathlib import Path +from typing import Literal, cast + +import tyro +from mjlab.envs import ManagerBasedRlEnvCfg, ManagerBasedRlEnv +from mjlab_husky.envs import G1SkaterManagerBasedRlEnvCfg, G1SkaterManagerBasedRlEnv +from mjlab_husky.rl import RslRlVecEnvWrapper +from mjlab_husky.rl import RslRlAMPOnPolicyRunnerCfg +from mjlab_husky.tasks.registry import list_tasks, load_env_cfg, load_rl_cfg, load_runner_cls +from mjlab.utils.gpu import select_gpus +from mjlab.utils.os import dump_yaml, get_checkpoint_path, get_wandb_checkpoint_path +from mjlab.utils.torch import configure_torch_backends +from mjlab.utils.wandb import add_wandb_tags +from mjlab.utils.wrappers import VideoRecorder +# os.environ["WANDB_MODE"] = "offline" + + +@dataclass(frozen=True) +class TrainConfig: + env: G1SkaterManagerBasedRlEnvCfg + agent: RslRlAMPOnPolicyRunnerCfg + registry_name: str | None = None + video: bool = False + video_length: int = 200 + video_interval: int = 2000 + enable_nan_guard: bool = False + torchrunx_log_dir: str | None = None + wandb_run_path: str | None = None + gpu_ids: list[int] | Literal["all"] | None = field(default_factory=lambda: [0]) + + @staticmethod + def from_task(task_id: str) -> "TrainConfig": + env_cfg = load_env_cfg(task_id) + agent_cfg = load_rl_cfg(task_id) + assert isinstance(agent_cfg, RslRlAMPOnPolicyRunnerCfg) + return TrainConfig(env=env_cfg, agent=agent_cfg) + + +def run_train(task_id: str, cfg: TrainConfig, log_dir: Path) -> None: + cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "") + if cuda_visible == "": + device = "cpu" + seed = cfg.agent.seed + rank = 0 + else: + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + rank = int(os.environ.get("RANK", "0")) + # Set EGL device to match the CUDA device. + os.environ["MUJOCO_EGL_DEVICE_ID"] = str(local_rank) + device = f"cuda:{local_rank}" + # Set seed to have diversity in different processes. + seed = cfg.agent.seed + local_rank + + configure_torch_backends() + + cfg.agent.seed = seed + cfg.env.seed = seed + + print(f"[INFO] Training with: device={device}, seed={seed}, rank={rank}") + + registry_name: str | None = None + + + # Enable NaN guard if requested. + if cfg.enable_nan_guard: + cfg.env.sim.nan_guard.enabled = True + print(f"[INFO] NaN guard enabled, output dir: {cfg.env.sim.nan_guard.output_dir}") + + if rank == 0: + print(f"[INFO] Logging experiment in directory: {log_dir}") + + env = G1SkaterManagerBasedRlEnv( + cfg=cfg.env, device=device, render_mode="rgb_array" if cfg.video else None + ) + + log_root_path = log_dir.parent # Go up from specific run dir to experiment dir. + + resume_path: Path | None = None + if cfg.agent.resume: + if cfg.wandb_run_path is not None: + # Load checkpoint from W&B. + resume_path, was_cached = get_wandb_checkpoint_path( + log_root_path, Path(cfg.wandb_run_path) + ) + if rank == 0: + run_id = resume_path.parent.name + checkpoint_name = resume_path.name + cached_str = "cached" if was_cached else "downloaded" + print( + f"[INFO]: Loading checkpoint from W&B: {checkpoint_name} " + f"(run: {run_id}, {cached_str})" + ) + else: + # Load checkpoint from local filesystem. + resume_path = get_checkpoint_path( + log_root_path, cfg.agent.load_run, cfg.agent.load_checkpoint + ) + + # Only record videos on rank 0 to avoid multiple workers writing to the same files. + if cfg.video and rank == 0: + env = VideoRecorder( + env, + video_folder=Path(log_dir) / "videos" / "train", + step_trigger=lambda step: step % cfg.video_interval == 0, + video_length=cfg.video_length, + disable_logger=True, + ) + print("[INFO] Recording videos during training.") + + env = RslRlVecEnvWrapper(env, clip_actions=cfg.agent.clip_actions) + + agent_cfg = asdict(cfg.agent) + env_cfg = asdict(cfg.env) + + runner_cls = load_runner_cls(task_id) + + runner_kwargs = {} + + + runner = runner_cls(env, agent_cfg, str(log_dir), device, **runner_kwargs) + + add_wandb_tags(cfg.agent.wandb_tags) + runner.add_git_repo_to_log(__file__) + if resume_path is not None: + print(f"[INFO]: Loading model checkpoint from: {resume_path}") + runner.load(str(resume_path)) + + # Only write config files from rank 0 to avoid race conditions. + if rank == 0: + dump_yaml(log_dir / "params" / "env.yaml", env_cfg) + dump_yaml(log_dir / "params" / "agent.yaml", agent_cfg) + + runner.learn( + num_learning_iterations=cfg.agent.max_iterations, init_at_random_ep_len=True + ) + + env.close() + + +def launch_training(task_id: str, args: TrainConfig | None = None): + args = args or TrainConfig.from_task(task_id) + + # Create log directory once before launching workers. + log_root_path = Path("logs") / "rsl_rl" / args.agent.experiment_name + log_root_path.resolve() + log_dir_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + if args.agent.run_name: + log_dir_name += f"_{args.agent.run_name}" + log_dir = log_root_path / log_dir_name + + # Select GPUs based on CUDA_VISIBLE_DEVICES and user specification. + selected_gpus, num_gpus = select_gpus(args.gpu_ids) + + # Set environment variables for all modes. + if selected_gpus is None: + os.environ["CUDA_VISIBLE_DEVICES"] = "" + else: + os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, selected_gpus)) + os.environ["MUJOCO_GL"] = "egl" + + if num_gpus <= 1: + # CPU or single GPU: run directly without torchrunx. + run_train(task_id, args, log_dir) + else: + # Multi-GPU: use torchrunx. + import torchrunx + + # torchrunx redirects stdout to logging. + logging.basicConfig(level=logging.INFO) + + # Configure torchrunx logging directory. + # Priority: 1) existing env var, 2) user flag, 3) default to {log_dir}/torchrunx. + if "TORCHRUNX_LOG_DIR" not in os.environ: + if args.torchrunx_log_dir is not None: + # User specified a value via flag (could be "" to disable). + os.environ["TORCHRUNX_LOG_DIR"] = args.torchrunx_log_dir + else: + # Default: put logs in training directory. + os.environ["TORCHRUNX_LOG_DIR"] = str(log_dir / "torchrunx") + + print(f"[INFO] Launching training with {num_gpus} GPUs", flush=True) + torchrunx.Launcher( + hostnames=["localhost"], + workers_per_host=num_gpus, + backend=None, # Let rsl_rl handle process group initialization. + copy_env_vars=torchrunx.DEFAULT_ENV_VARS_FOR_COPY + ("MUJOCO*",), + ).run(run_train, task_id, args, log_dir) + + +def main(): + # Parse first argument to choose the task. + # Import tasks to populate the registry. + all_tasks = list_tasks() + chosen_task, remaining_args = tyro.cli( + tyro.extras.literal_type_from_choices(all_tasks), + add_help=False, + return_unknown_args=True, + ) + + args = tyro.cli( + TrainConfig, + args=remaining_args, + default=TrainConfig.from_task(chosen_task), + prog=sys.argv[0] + f" {chosen_task}", + config=( + tyro.conf.AvoidSubcommands, + tyro.conf.FlagConversionOff, + ), + ) + del remaining_args + + launch_training(task_id=chosen_task, args=args) + + +if __name__ == "__main__": + main() diff --git a/src/mjlab_husky/tasks/__init__.py b/src/mjlab_husky/tasks/__init__.py new file mode 100644 index 0000000..26c873b --- /dev/null +++ b/src/mjlab_husky/tasks/__init__.py @@ -0,0 +1,5 @@ +from mjlab.utils.lab_api.tasks.importer import import_packages + +_BLACKLIST_PKGS = ["utils", ".mdp"] + +import_packages(__name__, _BLACKLIST_PKGS) diff --git a/src/mjlab_husky/tasks/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/tasks/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..4701c6e Binary files /dev/null and b/src/mjlab_husky/tasks/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/mjlab_husky/tasks/__pycache__/registry.cpython-312.pyc b/src/mjlab_husky/tasks/__pycache__/registry.cpython-312.pyc new file mode 100644 index 0000000..dd6e324 Binary files /dev/null and b/src/mjlab_husky/tasks/__pycache__/registry.cpython-312.pyc differ diff --git a/src/mjlab_husky/tasks/registry.py b/src/mjlab_husky/tasks/registry.py new file mode 100644 index 0000000..53e1b62 --- /dev/null +++ b/src/mjlab_husky/tasks/registry.py @@ -0,0 +1,71 @@ +"""Task registry system for managing environment registration and creation.""" + +from copy import deepcopy +from dataclasses import dataclass + +from mjlab.envs import ManagerBasedRlEnvCfg +from mjlab.rl import RslRlOnPolicyRunnerCfg + + +@dataclass +class _TaskCfg: + env_cfg: ManagerBasedRlEnvCfg + play_env_cfg: ManagerBasedRlEnvCfg + rl_cfg: RslRlOnPolicyRunnerCfg + runner_cls: type | None + + +# Private module-level registry: task_id -> task config. +_REGISTRY: dict[str, _TaskCfg] = {} + + +def register_mjlab_task( + task_id: str, + env_cfg: ManagerBasedRlEnvCfg, + play_env_cfg: ManagerBasedRlEnvCfg, + rl_cfg: RslRlOnPolicyRunnerCfg, + runner_cls: type | None = None, +) -> None: + """Register an environment task. + + Args: + task_id: Unique task identifier (e.g., "Mjlab-Velocity-Rough-Unitree-Go1"). + env_cfg: Environment configuration used for training. + play_env_cfg: Environment configuration in "play" mode. + rl_cfg: RL runner configuration. + runner_cls: Optional custom runner class. If None, uses OnPolicyRunner. + """ + if task_id in _REGISTRY: + raise ValueError(f"Task '{task_id}' is already registered") + _REGISTRY[task_id] = _TaskCfg(env_cfg, play_env_cfg, rl_cfg, runner_cls) + + +def list_tasks() -> list[str]: + """List all registered task IDs.""" + return sorted(_REGISTRY.keys()) + + +def load_env_cfg(task_name: str, play: bool = False) -> ManagerBasedRlEnvCfg: + """Load environment configuration for a task. + + Returns a deep copy to prevent mutation of the registered config. + """ + return deepcopy( + _REGISTRY[task_name].env_cfg if not play else _REGISTRY[task_name].play_env_cfg + ) + + +def load_rl_cfg(task_name: str) -> RslRlOnPolicyRunnerCfg: + """Load RL configuration for a task. + + Returns a deep copy to prevent mutation of the registered config. + """ + return deepcopy(_REGISTRY[task_name].rl_cfg) + + +def load_runner_cls(task_name: str) -> type | None: + """Load the runner class for a task. + + If None, the default OnPolicyRunner will be used. + """ + return _REGISTRY[task_name].runner_cls diff --git a/src/mjlab_husky/tasks/skater/__init__.py b/src/mjlab_husky/tasks/skater/__init__.py new file mode 100644 index 0000000..d78d6cc --- /dev/null +++ b/src/mjlab_husky/tasks/skater/__init__.py @@ -0,0 +1 @@ +"""Skating environments for legged robots.""" diff --git a/src/mjlab_husky/tasks/skater/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/tasks/skater/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..1ee1dc9 Binary files /dev/null and b/src/mjlab_husky/tasks/skater/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/mjlab_husky/tasks/skater/__pycache__/skater_env_cfg.cpython-312.pyc b/src/mjlab_husky/tasks/skater/__pycache__/skater_env_cfg.cpython-312.pyc new file mode 100644 index 0000000..1fe596a Binary files /dev/null and b/src/mjlab_husky/tasks/skater/__pycache__/skater_env_cfg.cpython-312.pyc differ diff --git a/src/mjlab_husky/tasks/skater/config/__init__.py b/src/mjlab_husky/tasks/skater/config/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mjlab_husky/tasks/skater/config/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/tasks/skater/config/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..d3cb0c4 Binary files /dev/null and b/src/mjlab_husky/tasks/skater/config/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/mjlab_husky/tasks/skater/config/g1/__init__.py b/src/mjlab_husky/tasks/skater/config/g1/__init__.py new file mode 100644 index 0000000..5567225 --- /dev/null +++ b/src/mjlab_husky/tasks/skater/config/g1/__init__.py @@ -0,0 +1,16 @@ +from mjlab_husky.tasks.registry import register_mjlab_task +from mjlab_husky.tasks.skater.rl import SkaterOnPolicyRunner + +from .env_cfgs import ( + unitree_g1_skater_env_cfg, +) +from .rl_cfg import unitree_g1_skater_ppo_runner_cfg + + +register_mjlab_task( + task_id="Mjlab-Skater-Flat-Unitree-G1", + env_cfg=unitree_g1_skater_env_cfg(), + play_env_cfg=unitree_g1_skater_env_cfg(play=True), + rl_cfg=unitree_g1_skater_ppo_runner_cfg(), + runner_cls=SkaterOnPolicyRunner, +) diff --git a/src/mjlab_husky/tasks/skater/config/g1/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/tasks/skater/config/g1/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..fed4296 Binary files /dev/null and b/src/mjlab_husky/tasks/skater/config/g1/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/mjlab_husky/tasks/skater/config/g1/__pycache__/env_cfgs.cpython-312.pyc b/src/mjlab_husky/tasks/skater/config/g1/__pycache__/env_cfgs.cpython-312.pyc new file mode 100644 index 0000000..622815d Binary files /dev/null and b/src/mjlab_husky/tasks/skater/config/g1/__pycache__/env_cfgs.cpython-312.pyc differ diff --git a/src/mjlab_husky/tasks/skater/config/g1/__pycache__/rl_cfg.cpython-312.pyc b/src/mjlab_husky/tasks/skater/config/g1/__pycache__/rl_cfg.cpython-312.pyc new file mode 100644 index 0000000..342849f Binary files /dev/null and b/src/mjlab_husky/tasks/skater/config/g1/__pycache__/rl_cfg.cpython-312.pyc differ diff --git a/src/mjlab_husky/tasks/skater/config/g1/env_cfgs.py b/src/mjlab_husky/tasks/skater/config/g1/env_cfgs.py new file mode 100644 index 0000000..1988540 --- /dev/null +++ b/src/mjlab_husky/tasks/skater/config/g1/env_cfgs.py @@ -0,0 +1,173 @@ +"""Unitree G1 skateboarding environment configurations.""" + +from mjlab_husky.asset_zoo.robots.skateboard.g1_skater_constants import ( + G1_23Dof_ACTION_SCALE, + get_g1_23dof_robot_cfg, + get_skateboard_cfg +) +from mjlab_husky.envs import G1SkaterManagerBasedRlEnvCfg +from mjlab.envs.mdp.actions import JointPositionActionCfg +from mjlab.managers.termination_manager import TerminationTermCfg +from mjlab.sensor import ContactMatch, ContactSensorCfg +from mjlab_husky.tasks.skater import mdp +from mjlab_husky.tasks.skater.mdp import SkateUniformVelocityCommandCfg +from mjlab_husky.tasks.skater.skater_env_cfg import make_g1_skater_env_cfg + + +def unitree_g1_skater_env_cfg(play: bool = False) -> G1SkaterManagerBasedRlEnvCfg: + cfg = make_g1_skater_env_cfg() + cfg.sim.njmax = 300 + cfg.sim.mujoco.ccd_iterations = 50 + cfg.sim.contact_sensor_maxmatch = 64 + cfg.sim.nconmax = 55 + + cfg.scene.entities = {"robot": get_g1_23dof_robot_cfg(), "skateboard": get_skateboard_cfg()} + + ######################################################### + ##### terrain ##### + ######################################################### + assert cfg.scene.terrain is not None + cfg.scene.terrain.terrain_type = "plane" + cfg.scene.terrain.terrain_generator = None + + ######################################################### + ##### contact sensors ##### + ######################################################### + left_feet_ground_cfg = ContactSensorCfg( + name="left_feet_ground_contact", + primary=ContactMatch( + mode="subtree", + pattern=r"^(left_ankle_roll_link)$", + entity="robot", + ), + secondary=ContactMatch(mode="body", pattern="terrain"), + fields=("found", "force"), + reduce="netforce", + num_slots=1, + track_air_time=True, + ) + + right_feet_ground_cfg = ContactSensorCfg( + name="right_feet_ground_contact", + primary=ContactMatch( + mode="subtree", + pattern=r"^(right_ankle_roll_link)$", + entity="robot", + ), + secondary=ContactMatch(mode="body", pattern="terrain"), + fields=("found", "force"), + reduce="netforce", + num_slots=1, + track_air_time=True, + ) + + left_feet_board_cfg = ContactSensorCfg( + name="left_feet_board_contact", + primary=ContactMatch( + mode="subtree", + pattern=r"^(left_ankle_roll_link)$", + entity="robot", + ), + secondary=ContactMatch(mode="geom", pattern="skateboard_marker_collision", entity="skateboard"), + fields=("found", "force"), + reduce="netforce", + num_slots=1, + track_air_time=True, + ) + + right_feet_board_cfg = ContactSensorCfg( + name="right_feet_board_contact", + primary=ContactMatch( + mode="subtree", + pattern=r"^(right_ankle_roll_link)$", + entity="robot", + ), + secondary=ContactMatch(mode="geom", pattern="skateboard_deck_collision", entity="skateboard"), + fields=("found", "force"), + reduce="netforce", + num_slots=1, + track_air_time=True, + ) + + robot_collision_cfg = ContactSensorCfg( + name="robot_collision", + primary=ContactMatch(mode="subtree", pattern="pelvis", entity="robot"), + secondary=ContactMatch(mode="subtree", pattern="pelvis", entity="robot"), + fields=("found",), + reduce="none", + num_slots=1, + ) + + skateboard_collision_cfg = ContactSensorCfg( + name="skateboard_collision", + primary=ContactMatch(mode="geom", pattern=r".*_wheel_collision$", entity="skateboard"), + secondary=ContactMatch(mode="body", pattern="terrain"), + fields=("found","force"), + reduce="none", + num_slots=1, + ) + + illegal_contact_cfg = ContactSensorCfg( + name="illegal_contact", + primary=ContactMatch(mode="geom", pattern=r".*_shin_collision|.*_linkage_brace_collision|.*_shoulder_yaw_collision|.*_elbow_yaw_collision|.*_wrist_collision|.*_hand_collision|pelvis_collision$", entity="robot"), + # secondary=ContactMatch(mode="body", pattern="terrain"), + fields=("found",), + reduce="none", + num_slots=1, + ) + + cfg.scene.sensors = (robot_collision_cfg,skateboard_collision_cfg,left_feet_ground_cfg, right_feet_ground_cfg,left_feet_board_cfg, right_feet_board_cfg,illegal_contact_cfg) + + + joint_pos_action = cfg.actions["joint_pos"] + assert isinstance(joint_pos_action, JointPositionActionCfg) + joint_pos_action.scale = G1_23Dof_ACTION_SCALE + + cfg.viewer.body_name = "torso_link" + + skate_cmd = cfg.commands["skate"] + assert isinstance(skate_cmd, SkateUniformVelocityCommandCfg) + skate_cmd.viz.z_offset = 1.15 + + + cfg.beizer_names = [ + "pelvis", + "left_hip_roll_link", + "left_knee_link", + "left_ankle_roll_link", + "right_hip_roll_link", + "right_knee_link", + "right_ankle_roll_link", + "torso_link", + "left_shoulder_roll_link", + "left_elbow_link", + "left_wrist_yaw_link", + "right_shoulder_roll_link", + "right_elbow_link", + "right_wrist_yaw_link", + ] + + cfg.slerp_names = cfg.beizer_names + cfg.phase_ratios = [0.0, 0.4, 0.5, 0.95, 1.0] + cfg.steer_init_pos = [ + -0.15, 0.1, 0.05, 0.6, -0.42, 0.0, + -0.15, -0.1, 0.05, 0.6, -0.42, 0.0, + 0, 0, 0.1, + 0, 0.55, -0.25, 0.55, + 0, -0.55, -0.25, 0.55 + ] + + # Apply play mode overrides. + if play: + # Effectively infinite episode length. + cfg.episode_length_s = int(60.0) + cfg.eval_mode = True + cfg.observations["policy"].enable_corruption = False + cfg.terminations = { + "time_out": TerminationTermCfg(func=mdp.time_out, time_out=True), + } + cfg.events.pop("push_robot", None) + # cfg.commands["skate"].ranges.lin_vel_x = (1.0, 1.0) # pyright: ignore[reportAttributeAccessIssue] + # uniform_(low, high) 要求 low <= high；固定航向偏角用 (a, a)，例如 (0.7, 0.7) 表示恒为 0.7 rad + cfg.commands["skate"].ranges.heading = (-0.7, 0) # pyright: ignore[reportAttributeAccessIssue] + return cfg \ No newline at end of file diff --git a/src/mjlab_husky/tasks/skater/config/g1/rl_cfg.py b/src/mjlab_husky/tasks/skater/config/g1/rl_cfg.py new file mode 100644 index 0000000..257546e --- /dev/null +++ b/src/mjlab_husky/tasks/skater/config/g1/rl_cfg.py @@ -0,0 +1,43 @@ +"""RL configuration for Unitree G1 skater task.""" + +from mjlab.rl import ( + RslRlPpoActorCriticCfg, + RslRlPpoAlgorithmCfg, +) + +from mjlab_husky.rl import ( + RslRlAMPOnPolicyRunnerCfg +) + + +def unitree_g1_skater_ppo_runner_cfg() -> RslRlAMPOnPolicyRunnerCfg: + """Create RL runner configuration for Unitree G1 skater task.""" + return RslRlAMPOnPolicyRunnerCfg( + policy=RslRlPpoActorCriticCfg( + init_noise_std=1.0, + actor_obs_normalization=True, + critic_obs_normalization=True, + actor_hidden_dims=(512, 256, 128), + critic_hidden_dims=(512, 256, 128), + activation="elu", + ), + algorithm=RslRlPpoAlgorithmCfg( + value_loss_coef=1.0, + use_clipped_value_loss=True, + clip_param=0.2, + entropy_coef=0.005, + num_learning_epochs=5, + num_mini_batches=4, + learning_rate=1.0e-3, + schedule="adaptive", + gamma=0.99, + lam=0.95, + desired_kl=0.01, + max_grad_norm=1.0, + class_name="AMP_PPO", + ), + experiment_name="g1_skater", + save_interval=500, + num_steps_per_env=24, + max_iterations=50_000, + ) diff --git a/src/mjlab_husky/tasks/skater/mdp/__init__.py b/src/mjlab_husky/tasks/skater/mdp/__init__.py new file mode 100644 index 0000000..be53f71 --- /dev/null +++ b/src/mjlab_husky/tasks/skater/mdp/__init__.py @@ -0,0 +1,6 @@ +from mjlab.envs.mdp import * # noqa: F401, F403 + +from .observations import * # noqa: F403 +from .rewards import * # noqa: F403 +from .terminations import * # noqa: F403 +from .velocity_command import * # noqa: F403 diff --git a/src/mjlab_husky/tasks/skater/mdp/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/tasks/skater/mdp/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..a419060 Binary files /dev/null and b/src/mjlab_husky/tasks/skater/mdp/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/mjlab_husky/tasks/skater/mdp/__pycache__/observations.cpython-312.pyc b/src/mjlab_husky/tasks/skater/mdp/__pycache__/observations.cpython-312.pyc new file mode 100644 index 0000000..c1d816f Binary files /dev/null and b/src/mjlab_husky/tasks/skater/mdp/__pycache__/observations.cpython-312.pyc differ diff --git a/src/mjlab_husky/tasks/skater/mdp/__pycache__/rewards.cpython-312.pyc b/src/mjlab_husky/tasks/skater/mdp/__pycache__/rewards.cpython-312.pyc new file mode 100644 index 0000000..2348837 Binary files /dev/null and b/src/mjlab_husky/tasks/skater/mdp/__pycache__/rewards.cpython-312.pyc differ diff --git a/src/mjlab_husky/tasks/skater/mdp/__pycache__/terminations.cpython-312.pyc b/src/mjlab_husky/tasks/skater/mdp/__pycache__/terminations.cpython-312.pyc new file mode 100644 index 0000000..e11c38b Binary files /dev/null and b/src/mjlab_husky/tasks/skater/mdp/__pycache__/terminations.cpython-312.pyc differ diff --git a/src/mjlab_husky/tasks/skater/mdp/__pycache__/velocity_command.cpython-312.pyc b/src/mjlab_husky/tasks/skater/mdp/__pycache__/velocity_command.cpython-312.pyc new file mode 100644 index 0000000..bf7b98d Binary files /dev/null and b/src/mjlab_husky/tasks/skater/mdp/__pycache__/velocity_command.cpython-312.pyc differ diff --git a/src/mjlab_husky/tasks/skater/mdp/observations.py b/src/mjlab_husky/tasks/skater/mdp/observations.py new file mode 100644 index 0000000..2152dc6 --- /dev/null +++ b/src/mjlab_husky/tasks/skater/mdp/observations.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from mjlab.managers.scene_entity_config import SceneEntityCfg +from mjlab.sensor import ContactSensor + +if TYPE_CHECKING: + from mjlab_husky.envs import G1SkaterManagerBasedRlEnv + +_DEFAULT_ASSET_CFG = SceneEntityCfg("robot") +from mjlab.utils.lab_api.math import ( + quat_apply, + quat_mul, + quat_apply_inverse, + quat_conjugate, + wrap_to_pi, +) + + +def heading(env: G1SkaterManagerBasedRlEnv) -> torch.Tensor: + return env.robot.data.heading_w.unsqueeze(-1) + +def heading_error(env: G1SkaterManagerBasedRlEnv, command_name: str) -> torch.Tensor: + target_w = env.get_heading_target_w(command_name) + heading_w = env.skateboard.data.heading_w + if target_w is not None: + error = wrap_to_pi(target_w - heading_w) + else: + command = env.command_manager.get_command(command_name) + assert command is not None + error = wrap_to_pi(command[:, 1] - heading_w) + return error.unsqueeze(-1) + +def contact_phase(env: G1SkaterManagerBasedRlEnv) -> torch.Tensor: + return env.contact_phase.clone() + +def phase(env: G1SkaterManagerBasedRlEnv) -> torch.Tensor: + phase = env._get_phase().clone() + return phase.unsqueeze(1) + +def skate_pose_local(env: G1SkaterManagerBasedRlEnv) -> torch.Tensor: + skate_pos = env.skateboard.data.root_link_pos_w.squeeze(1) - env.robot.data.root_link_pos_w + skate_pos_local = quat_apply_inverse(env.robot.data.root_link_quat_w, skate_pos) + + skate_quat_local = quat_mul(quat_conjugate(env.robot.data.root_link_quat_w), env.skateboard.data.root_link_quat_w.squeeze(1)) + skate_rot_6d_local = quaternion_to_tangent_and_normal(skate_quat_local) + + skate_states = torch.cat([skate_pos_local, skate_rot_6d_local], dim=-1) + return skate_states + +def skate_vel_local(env: G1SkaterManagerBasedRlEnv) -> torch.Tensor: + skate_vel = env.skateboard.data.root_link_lin_vel_w + skate_vel_local = quat_apply_inverse(env.robot.data.root_link_quat_w, skate_vel) + return skate_vel_local + +def skate_ang_vel_local(env: G1SkaterManagerBasedRlEnv) -> torch.Tensor: + skate_ang_vel = env.skateboard.data.root_link_ang_vel_w + skate_ang_vel_local = quat_apply_inverse(env.robot.data.root_link_quat_w, skate_ang_vel) + return skate_ang_vel_local + +def skateboard_roll(env: G1SkaterManagerBasedRlEnv,) -> torch.Tensor: + return env.skateboard.data.joint_pos[:, [0, 1, 4]].view(env.num_envs, -1) + +def trans_target_pos_b(env: G1SkaterManagerBasedRlEnv,) -> torch.Tensor: + target_pos_b, _, _ = env._get_transition_target_b() + return target_pos_b[:, env.beizer_ids, :].view(env.num_envs, -1) + +def trans_target_quat_b(env: G1SkaterManagerBasedRlEnv,) -> torch.Tensor: + _, target_quat_b, _ = env._get_transition_target_b() + return target_quat_b[:, env.slerp_ids, :].view(env.num_envs, -1) + +def foot_contact_forces(env: G1SkaterManagerBasedRlEnv, sensor_name: str) -> torch.Tensor: + sensor: ContactSensor = env.scene[sensor_name] + sensor_data = sensor.data + assert sensor_data.force is not None + forces_flat = sensor_data.force.flatten(start_dim=1) # [B, N*3] + return torch.sign(forces_flat) * torch.log1p(torch.abs(forces_flat)) + +@torch.jit.script +def quaternion_to_tangent_and_normal(q: torch.Tensor) -> torch.Tensor: + ref_tangent = torch.zeros_like(q[..., :3]) + ref_normal = torch.zeros_like(q[..., :3]) + ref_tangent[..., 0] = 1 + ref_normal[..., -1] = 1 + tangent = quat_apply(q, ref_tangent) + normal = quat_apply(q, ref_normal) + return torch.cat([tangent, normal], dim=len(tangent.shape) - 1) + + diff --git a/src/mjlab_husky/tasks/skater/mdp/rewards.py b/src/mjlab_husky/tasks/skater/mdp/rewards.py new file mode 100644 index 0000000..85dbf28 --- /dev/null +++ b/src/mjlab_husky/tasks/skater/mdp/rewards.py @@ -0,0 +1,199 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from mjlab.entity import Entity +from mjlab.managers.scene_entity_config import SceneEntityCfg +from mjlab.sensor import ContactSensor + +from mjlab.utils.lab_api.math import ( + wrap_to_pi, + quat_mul, + quat_apply_inverse, + quat_error_magnitude, + euler_xyz_from_quat +) +if TYPE_CHECKING: + from mjlab_husky.envs import G1SkaterManagerBasedRlEnv + + +_DEFAULT_ASSET_CFG = SceneEntityCfg("robot") + + +#### push phase rewards #### +def push_skateboard_lin_vel( + env: G1SkaterManagerBasedRlEnv, + std: float, + command_name: str, + asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG, +) -> torch.Tensor: + asset: Entity = env.scene[asset_cfg.name] + command = env.command_manager.get_command(command_name) + assert command is not None, f"Command '{command_name}' not found." + actual = asset.data.root_link_lin_vel_b + x_error = torch.square(command[:, 0] - actual[:, 0]) + y_error = torch.square(actual[:, 1]) + z_error = torch.square(actual[:, 2]) + lin_vel_error = x_error + y_error + z_error + return torch.exp(-lin_vel_error / std**2) + +def push_yaw_align(env: G1SkaterManagerBasedRlEnv, std: float) -> torch.Tensor: + _, _, yaw = euler_xyz_from_quat(env.robot.data.root_link_quat_w) + _, _, skateboard_yaw = euler_xyz_from_quat(env.skateboard.data.root_link_quat_w.squeeze(1)) + yaw_diff = torch.abs(yaw - skateboard_yaw) + return torch.exp(-yaw_diff / std**2) + +def feet_air_time( + env: G1SkaterManagerBasedRlEnv, + sensor_name: str, + threshold_min: float, + threshold_max: float, + command_name: str, + command_threshold: float, +) -> torch.Tensor: + """Reward feet air time.""" + sensor: ContactSensor = env.scene[sensor_name] + sensor_data = sensor.data + current_air_time = sensor_data.current_air_time + assert current_air_time is not None + in_range = (current_air_time > threshold_min) & (current_air_time < threshold_max) + reward = torch.sum(in_range.float(), dim=1) + in_air = current_air_time > 0 + num_in_air = torch.sum(in_air.float()) + push_envs = env.contact_phase[:,0] == 1. + mean_air_time = torch.sum(current_air_time[push_envs] * in_air[push_envs].float()) / torch.clamp( + num_in_air, min=1 + ) + env.extras["log"]["Metrics/air_time_mean"] = mean_air_time + command = env.command_manager.get_command(command_name) + + assert command is not None + scale = (command[:, 0] > command_threshold).float() + reward *= scale + return reward + +def push_contact_ground_parallel(env: G1SkaterManagerBasedRlEnv) -> torch.Tensor: + left_ankle_pos = env.robot.data.site_pos_w[:, env.left_foot_site_ids, 2].clone() * 10 + var = left_ankle_pos.var(1) + var = torch.mean(var.view(-1, 1), dim=-1) + reward = var < 0.05 + left_feet_ground_contact = env._get_feet_contact_g()[:, 0] + reward = reward * left_feet_ground_contact.float() + return reward + +#### steer phase rewards #### +def steer_contact_num(env: G1SkaterManagerBasedRlEnv) -> torch.Tensor: + feet_contact_b = env._get_feet_contact_b() + both_contact = torch.sum(feet_contact_b, dim=-1) == 2 + feet_contact_g = env._get_feet_contact_g() + left_ground_contact = feet_contact_g[:,0] + return 2* both_contact.float() - left_ground_contact.float() + +def steer_joint_pos(env: G1SkaterManagerBasedRlEnv, std: float) -> torch.Tensor: + dof_error = torch.mean(torch.square(env.robot.data.joint_pos - env.steer_init_pos), dim=1) + return torch.exp(-dof_error / std**2) + +def steer_feet_dis(env: G1SkaterManagerBasedRlEnv, std: float) -> torch.Tensor: + dis = env._get_feet_marker_dis() + skateb_contact_dis = torch.norm(dis, dim=-1).mean(dim=-1) + reward = torch.exp(-skateb_contact_dis / std**2) + return reward + +def steer_track_heading(env: G1SkaterManagerBasedRlEnv, command_name: str, std: float) -> torch.Tensor: + target_w = env.get_heading_target_w(command_name) + heading_w = env.skateboard.data.heading_w + if target_w is not None: + error = wrap_to_pi(heading_w - target_w) + else: + command = env.command_manager.get_command(command_name) + assert command is not None + error = wrap_to_pi(heading_w - command[:, 1]) + r = torch.exp(-torch.abs(error) / (std**2)) + in_steer = env.contact_phase[:, 1] == 1.0 + reward = torch.where(in_steer, r, torch.zeros_like(r)) + return reward + +def steer_tilt_guide(env: G1SkaterManagerBasedRlEnv, command_name: str, std: float) -> torch.Tensor: + gamma = env.skateboard.data.joint_pos[:, 0] + target_w = env.get_heading_target_w(command_name) + heading_w = env.skateboard.data.heading_w + if target_w is not None: + delta_theta = wrap_to_pi(target_w - heading_w) + else: + command = env.command_manager.get_command(command_name) + assert command is not None + delta_theta = wrap_to_pi(command[:, 1] - heading_w) + vx = env.skateboard.data.root_link_lin_vel_b[:, 0] + remaining_steps = env._steer_remaining_steps() + delta_t = (remaining_steps * env.step_dt).clamp(min=0.5) + lam = torch.deg2rad(torch.tensor(env.cfg.rake_angle, device=env.device, dtype=torch.float32)) + tan_sigma = (0.4 * delta_theta) / (vx * delta_t + 1e-6) + sin_gamma = torch.clamp(tan_sigma / torch.tan(lam), -0.99, 0.99) + gamma_ref = torch.clip(torch.asin(sin_gamma), -0.2, 0.2) + steer_envs = env.contact_phase[:, 1] == 1.0 + diff_gamma = torch.abs(gamma - gamma_ref) + reward = torch.exp(-diff_gamma / std**2) + reward = torch.where(steer_envs, reward, torch.zeros_like(reward)) + return reward + +#### transition rewards #### +def transition_body_pos_tracking(env: G1SkaterManagerBasedRlEnv, std: float) -> torch.Tensor: + target_pos_b, _, in_transition = env._get_transition_target_b() + + body_pos_w = env.robot.data.body_link_pos_w[:, :, :3] + root_pos_w = env.skateboard.data.root_link_pos_w[:, :3][:, None, :].repeat(1, env.robot.num_bodies, 1) + root_quat_w = env.skateboard.data.root_link_quat_w[:, None, :].repeat(1, env.robot.num_bodies, 1) + + rel_pos_w = body_pos_w - root_pos_w + current_body_pos_b = quat_apply_inverse(root_quat_w, rel_pos_w) + + pos_error = (current_body_pos_b - target_pos_b)[:, env.beizer_ids, :] + pos_error_norm = torch.sum(torch.square(pos_error),dim=-1) + reward = torch.exp(- pos_error_norm.mean(dim=-1) / std**2) + reward = torch.where(in_transition, reward, torch.zeros_like(reward)) + return reward + +def transition_body_rot_tracking(env: G1SkaterManagerBasedRlEnv, std: float) -> torch.Tensor: + _, target_quat_b, in_transition = env._get_transition_target_b() + + body_quat_w = env.robot.data.body_link_quat_w[:, :, :4] + root_quat_w = env.skateboard.data.root_link_quat_w[:, None, :].repeat(1, env.robot.num_bodies, 1) + target_quat_w = quat_mul(root_quat_w, target_quat_b) + + quat_error = torch.square(quat_error_magnitude(target_quat_w[:, env.slerp_ids, :],body_quat_w[:, env.slerp_ids, :])) + reward = torch.exp(- quat_error.mean(dim=-1) / std**2) + reward = torch.where(in_transition, reward, torch.zeros_like(reward)) + return reward + +def transition_penalty_contact(env: G1SkaterManagerBasedRlEnv, sensor_name: str) -> torch.Tensor: + sensor: ContactSensor = env.scene[sensor_name] + assert sensor.data.found is not None + return sensor.data.found.squeeze(-1) + +#### regularization rewards #### +def reg_wheel_contact_number(env: G1SkaterManagerBasedRlEnv) -> torch.Tensor: + wheel_contact_number = torch.sum(env.wheel_contact_filt, dim=1) + reward = wheel_contact_number == 4 + return reward + +def self_collision_cost(env: G1SkaterManagerBasedRlEnv, sensor_name: str) -> torch.Tensor: + """Cost that returns the number of self-collisions detected by a sensor.""" + sensor: ContactSensor = env.scene[sensor_name] + assert sensor.data.found is not None + return sensor.data.found.squeeze(-1) + +def stand_still(env: G1SkaterManagerBasedRlEnv, std: float) -> torch.Tensor: + still_envs = env.still.clone() + dof_error = torch.mean(torch.square(env.robot.data.joint_pos - env.robot.data.default_joint_pos),dim=1) + reward = torch.exp(-dof_error / std**2) + return reward * still_envs.float() + +def board_flat(env: G1SkaterManagerBasedRlEnv, std: float) -> torch.Tensor: + gamma = env.skateboard.data.joint_pos[:, 0] + non_steer = env.contact_phase[:, 1] != 1.0 + diff_gamma = torch.abs(gamma) + reward = torch.exp(-diff_gamma / std**2) + reward = torch.where(non_steer, reward, torch.zeros_like(reward)) + return reward + \ No newline at end of file diff --git a/src/mjlab_husky/tasks/skater/mdp/terminations.py b/src/mjlab_husky/tasks/skater/mdp/terminations.py new file mode 100644 index 0000000..ebd1847 --- /dev/null +++ b/src/mjlab_husky/tasks/skater/mdp/terminations.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from mjlab.sensor import ContactSensor + +if TYPE_CHECKING: + from mjlab_husky.envs import G1SkaterManagerBasedRlEnv + + +def illegal_contact(env: G1SkaterManagerBasedRlEnv, sensor_name: str) -> torch.Tensor: + sensor: ContactSensor = env.scene[sensor_name] + assert sensor.data.found is not None + return torch.any(sensor.data.found, dim=-1) + + +def bad_feet_off_board(env: G1SkaterManagerBasedRlEnv) -> torch.Tensor: + feet_contact_b = env._get_feet_contact_b() + bad_contact = torch.sum(feet_contact_b, dim=-1) == 0 + return bad_contact diff --git a/src/mjlab_husky/tasks/skater/mdp/velocity_command.py b/src/mjlab_husky/tasks/skater/mdp/velocity_command.py new file mode 100644 index 0000000..b05f193 --- /dev/null +++ b/src/mjlab_husky/tasks/skater/mdp/velocity_command.py @@ -0,0 +1,163 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +import numpy as np +import torch + +from mjlab.managers.command_manager import CommandTerm, CommandTermCfg +from mjlab.utils.lab_api.math import ( + wrap_to_pi, + quat_apply_inverse, + yaw_quat, +) + +if TYPE_CHECKING: + from mjlab_husky.envs import G1SkaterManagerBasedRlEnv + from mjlab.viewer.debug_visualizer import DebugVisualizer + +@dataclass(kw_only=True) +class SkateUniformVelocityCommandCfg(CommandTermCfg): + heading_command: bool = False + rel_standing_envs: float = 0.0 + rel_heading_envs: float = 1.0 + @dataclass + class Ranges: + lin_vel_x: tuple[float, float] + heading: tuple[float, float] | None = None + + ranges: Ranges + + @dataclass + class VizCfg: + z_offset: float = 0.2 + scale: float = 0.5 + heading_arrow_length: float = 1.25 + """World-space arrow length (m). Keep near character scale (~1–1.5); very long + arrows look like a thin needle with a tiny head in the viewer.""" + heading_arrow_width: float | None = None + """Shaft width; if None, uses meansize * heading_arrow_width_frac (clamped).""" + heading_arrow_width_frac: float = 0.28 + + viz: VizCfg = field(default_factory=VizCfg) + + def build(self, env: G1SkaterManagerBasedRlEnv) -> SkateUniformVelocityCommand: + return SkateUniformVelocityCommand(self, env) + + def __post_init__(self): + if self.heading_command and self.ranges.heading is None: + raise ValueError( + "The velocity command has heading commands active (heading_command=True) but " + "the `ranges.heading` parameter is set to None." + ) + +class SkateUniformVelocityCommand(CommandTerm): + cfg: SkateUniformVelocityCommandCfg + + def __init__(self, cfg: SkateUniformVelocityCommandCfg, env: G1SkaterManagerBasedRlEnv): + super().__init__(cfg, env) + if self.cfg.heading_command and self.cfg.ranges.heading is None: + raise ValueError("heading_command=True but ranges.heading is set to None.") + if self.cfg.ranges.heading and not self.cfg.heading_command: + raise ValueError("ranges.heading is set but heading_command=False.") + + self.robot = env.robot + self.skateboard = env.skateboard + self.env = env + + self.command_b = torch.zeros(self.num_envs, 2, device=self.device) + self.heading_ref_b = torch.zeros(self.num_envs, device=self.device) + self.heading_error = torch.zeros(self.num_envs, device=self.device) + self.is_heading_env = torch.zeros( + self.num_envs, dtype=torch.bool, device=self.device + ) + self.is_standing_env = torch.zeros_like(self.is_heading_env) + + self.metrics["error_vel"] = torch.zeros(self.num_envs, device=self.device) + self.metrics["error_yaw"] = torch.zeros(self.num_envs, device=self.device) + + @property + def command(self) -> torch.Tensor: + return self.command_b + + @property + def target_heading_w(self) -> torch.Tensor: + return wrap_to_pi(self.heading_ref_b + self.command_b[:, 1]) + + def _update_metrics(self): + # time for which the command was executed + max_command_time = self.cfg.resampling_time_range[1] + max_command_step = max_command_time / self._env.step_dt + # logs data + vel_yaw = quat_apply_inverse( + yaw_quat(self.skateboard.data.root_link_quat_w), + self.skateboard.data.root_link_lin_vel_w) + self.metrics["error_vel"] += ( + torch.abs(self.command_b[:, 0] - vel_yaw[:, 0]) / max_command_step + ) * self.env.contact_phase[:, 0] + self.metrics["error_yaw"] += ( + torch.abs(wrap_to_pi(self.target_heading_w - self.skateboard.data.heading_w)) / max_command_step + ) * self.env.contact_phase[:, 1] + + def _resample_command(self, env_ids: torch.Tensor) -> None: + r = torch.empty(len(env_ids), device=self.device) + self.command_b[env_ids, 0] = r.uniform_(*self.cfg.ranges.lin_vel_x) + if self.cfg.heading_command: + assert self.cfg.ranges.heading is not None + self.is_heading_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_heading_envs + heading_resample_ids = env_ids[self.is_heading_env[env_ids]] + if len(heading_resample_ids) > 0: + self._resample_heading_command(heading_resample_ids) + self.is_standing_env[env_ids] = r.uniform_(0.0, 1.0) <= self.cfg.rel_standing_envs + + def _update_command(self): + if self.cfg.heading_command: + env_ids = self.is_heading_env.nonzero(as_tuple=False).flatten() + if len(env_ids) > 0: + resample_mask = self.env.just_exited_push2steer + resample_env_ids = env_ids[resample_mask] + if len(resample_env_ids) > 0: + self._resample_heading_command(resample_env_ids) + standing_env_ids = self.is_standing_env.nonzero(as_tuple=False).flatten() + self.command_b[standing_env_ids] = 0.0 + + def _resample_heading_command(self, env_ids: torch.Tensor) -> None: + if len(env_ids) == 0: + return + if not self.cfg.heading_command or self.cfg.ranges.heading is None: + return + r = torch.empty(len(env_ids), device=self.device) + current_heading = self.skateboard.data.heading_w[env_ids] + relative_heading = r.uniform_(*self.cfg.ranges.heading) + self.command_b[env_ids, 1] = relative_heading + self.heading_ref_b[env_ids] = current_heading + + def _debug_vis_impl(self, visualizer: "DebugVisualizer") -> None: + super()._debug_vis_impl(visualizer) + + batch = visualizer.env_idx + if batch >= self.num_envs: + return + if not self.cfg.heading_command or not self.is_heading_env[batch]: + return + + skate_pos = self.skateboard.data.root_link_pos_w[batch].cpu().numpy() + yaw = self.target_heading_w[batch].item() + viz = self.cfg.viz + heading_goal_distance = viz.heading_arrow_length + direction = np.array([np.cos(yaw), np.sin(yaw), 0.0]) + goal_pos = skate_pos + heading_goal_distance * direction + ms = float(visualizer.meansize) + if viz.heading_arrow_width is not None: + arrow_w = float(viz.heading_arrow_width) + else: + arrow_w = max(0.03, ms * viz.heading_arrow_width_frac) + if self.env.contact_phase[batch, 1] > 0.5: + visualizer.add_arrow( + start=skate_pos, + end=goal_pos, + color=(1.0, 0.0, 0.0, 1.0), + width=arrow_w, + label="heading_target", + ) \ No newline at end of file diff --git a/src/mjlab_husky/tasks/skater/rl/__init__.py b/src/mjlab_husky/tasks/skater/rl/__init__.py new file mode 100644 index 0000000..151660b --- /dev/null +++ b/src/mjlab_husky/tasks/skater/rl/__init__.py @@ -0,0 +1,3 @@ +from mjlab_husky.tasks.skater.rl.runner import ( + SkaterOnPolicyRunner as SkaterOnPolicyRunner, +) diff --git a/src/mjlab_husky/tasks/skater/rl/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/tasks/skater/rl/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..58ce4ce Binary files /dev/null and b/src/mjlab_husky/tasks/skater/rl/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/mjlab_husky/tasks/skater/rl/__pycache__/exporter.cpython-312.pyc b/src/mjlab_husky/tasks/skater/rl/__pycache__/exporter.cpython-312.pyc new file mode 100644 index 0000000..91e4876 Binary files /dev/null and b/src/mjlab_husky/tasks/skater/rl/__pycache__/exporter.cpython-312.pyc differ diff --git a/src/mjlab_husky/tasks/skater/rl/__pycache__/runner.cpython-312.pyc b/src/mjlab_husky/tasks/skater/rl/__pycache__/runner.cpython-312.pyc new file mode 100644 index 0000000..ec29b82 Binary files /dev/null and b/src/mjlab_husky/tasks/skater/rl/__pycache__/runner.cpython-312.pyc differ diff --git a/src/mjlab_husky/tasks/skater/rl/exporter.py b/src/mjlab_husky/tasks/skater/rl/exporter.py new file mode 100644 index 0000000..93626b7 --- /dev/null +++ b/src/mjlab_husky/tasks/skater/rl/exporter.py @@ -0,0 +1,37 @@ +import os + +from mjlab.envs import ManagerBasedRlEnv +from mjlab.rl.exporter_utils import ( + attach_metadata_to_onnx, + get_base_metadata, +) +from mjlab.utils.lab_api.rl.exporter import _OnnxPolicyExporter + + +def export_skater_policy_as_onnx( + actor_critic: object, + path: str, + normalizer: object | None = None, + filename="policy.onnx", + verbose=False, +): + if not os.path.exists(path): + os.makedirs(path, exist_ok=True) + policy_exporter = _OnnxPolicyExporter(actor_critic, normalizer, verbose) + policy_exporter.export(path, filename) + + +def attach_onnx_metadata( + env: ManagerBasedRlEnv, run_path: str, path: str, filename="policy.onnx" +) -> None: + """Attach skater-specific metadata to ONNX model. + + Args: + env: The RL environment. + run_path: W&B run path or other identifier. + path: Directory containing the ONNX file. + filename: Name of the ONNX file. + """ + onnx_path = os.path.join(path, filename) + metadata = get_base_metadata(env, run_path) # Skater has no extra metadata. + attach_metadata_to_onnx(onnx_path, metadata) diff --git a/src/mjlab_husky/tasks/skater/rl/runner.py b/src/mjlab_husky/tasks/skater/rl/runner.py new file mode 100644 index 0000000..95e0da7 --- /dev/null +++ b/src/mjlab_husky/tasks/skater/rl/runner.py @@ -0,0 +1,37 @@ +import os + +import wandb +from rsl_rl.runners import AMPOnPolicyRunner +from mjlab.rl import RslRlVecEnvWrapper +from mjlab_husky.tasks.skater.rl.exporter import ( + attach_onnx_metadata, + export_skater_policy_as_onnx, +) + + +class SkaterOnPolicyRunner(AMPOnPolicyRunner): + env: RslRlVecEnvWrapper + + def save(self, path: str, infos=None): + """Save the model and training information.""" + super().save(path, infos) + if self.logger_type in ["wandb"]: + policy_path = path.split("model")[0] + filename = os.path.basename(os.path.dirname(policy_path)) + ".onnx" + if self.alg.policy.actor_obs_normalization: + normalizer = self.alg.policy.actor_obs_normalizer + else: + normalizer = None + export_skater_policy_as_onnx( + self.alg.policy, + normalizer=normalizer, + path=policy_path, + filename=filename, + ) + attach_onnx_metadata( + self.env.unwrapped, + wandb.run.name, # type: ignore + path=policy_path, + filename=filename, + ) + wandb.save(policy_path + filename, base_path=os.path.dirname(policy_path)) diff --git a/src/mjlab_husky/tasks/skater/skater_env_cfg.py b/src/mjlab_husky/tasks/skater/skater_env_cfg.py new file mode 100644 index 0000000..f0644ca --- /dev/null +++ b/src/mjlab_husky/tasks/skater/skater_env_cfg.py @@ -0,0 +1,389 @@ +import math +from mjlab_husky.envs import G1SkaterManagerBasedRlEnvCfg +from mjlab.envs.mdp.actions import JointPositionActionCfg +from mjlab.managers.action_manager import ActionTermCfg +from mjlab.managers.command_manager import CommandTermCfg +from mjlab.managers.event_manager import EventTermCfg +from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg +from mjlab.managers.reward_manager import RewardTermCfg +from mjlab.managers.scene_entity_config import SceneEntityCfg +from mjlab.managers.termination_manager import TerminationTermCfg +from mjlab.scene import SceneCfg +from mjlab.sim import MujocoCfg, SimulationCfg +from mjlab_husky.tasks.skater import mdp +from mjlab_husky.tasks.skater.mdp import SkateUniformVelocityCommandCfg +from mjlab.terrains import TerrainImporterCfg +from mjlab.utils.noise import UniformNoiseCfg as Unoise +from mjlab.viewer import ViewerConfig + + +def make_g1_skater_env_cfg() -> G1SkaterManagerBasedRlEnvCfg: + ## + # Observations + ## + + policy_terms = { + "command": ObservationTermCfg( + func=mdp.generated_commands, + params={"command_name": "skate"}, + scale=(2.0, 1.0), + ), + "heading": ObservationTermCfg( + func=mdp.heading, + scale=1.0 / math.pi, + ), + "base_ang_vel": ObservationTermCfg( + func=mdp.builtin_sensor, + params={"sensor_name": "robot/imu_ang_vel"}, + noise=Unoise(n_min=-0.2, n_max=0.2), + scale=0.25, + ), + "projected_gravity": ObservationTermCfg( + func=mdp.projected_gravity, + noise=Unoise(n_min=-0.05, n_max=0.05), + ), + "joint_pos": ObservationTermCfg( + func=mdp.joint_pos_rel, + noise=Unoise(n_min=-0.01, n_max=0.01), + ), + "joint_vel": ObservationTermCfg( + func=mdp.joint_vel_rel, + noise=Unoise(n_min=-1.5, n_max=1.5), + scale=0.05, + ), + "actions": ObservationTermCfg(func=mdp.last_action), + "phase": ObservationTermCfg(func=mdp.phase), + } + + critic_terms = { + **policy_terms, + "base_lin_vel": ObservationTermCfg( + func=mdp.builtin_sensor, + params={"sensor_name": "robot/imu_lin_vel"}, + ), + "heading_error": ObservationTermCfg( + func=mdp.heading_error, + params={"command_name": "skate"}, + scale=1.0 / math.pi, + ), + "skate_pose_local": ObservationTermCfg(func=mdp.skate_pose_local), + "skate_vel_local": ObservationTermCfg(func=mdp.skate_vel_local), + "skate_ang_vel_local": ObservationTermCfg(func=mdp.skate_ang_vel_local), + "skateboard_roll": ObservationTermCfg(func=mdp.skateboard_roll), + "trans_target_pos_b": ObservationTermCfg(func=mdp.trans_target_pos_b), + "trans_target_quat_b": ObservationTermCfg(func=mdp.trans_target_quat_b), + "l_foot_contact_forces": ObservationTermCfg( + func=mdp.foot_contact_forces, + params={"sensor_name": "left_feet_ground_contact"}, + ), + "r_foot_contact_forces": ObservationTermCfg( + func=mdp.foot_contact_forces, + params={"sensor_name": "right_feet_ground_contact"}, + ), + "l_foot_contact_forces_b": ObservationTermCfg( + func=mdp.foot_contact_forces, + params={"sensor_name": "left_feet_board_contact"}, + ), + "r_foot_contact_forces_b": ObservationTermCfg( + func=mdp.foot_contact_forces, + params={"sensor_name": "right_feet_board_contact"}, + ), + "contact_phase": ObservationTermCfg(func=mdp.contact_phase), + } + + observations = { + "policy": ObservationGroupCfg( + terms=policy_terms, + concatenate_terms=True, + enable_corruption=True, + history_length=5, + flatten_history_dim=True, + ), + "critic": ObservationGroupCfg( + terms=critic_terms, + concatenate_terms=True, + enable_corruption=False, + ), + } + + ## + # Actions + ## + + actions: dict[str, ActionTermCfg] = { + "joint_pos": JointPositionActionCfg( + entity_name="robot", + actuator_names=(".*",), + scale=0.5, # Override per-robot. + use_default_offset=True, + ) + } + + + ## + # Commands + ## + + commands: dict[str, CommandTermCfg] = { + "skate": SkateUniformVelocityCommandCfg( + resampling_time_range=(20.0, 20.0), + rel_standing_envs=0.0, + rel_heading_envs=1.0, + heading_command=True, + debug_vis=True, + ranges=SkateUniformVelocityCommandCfg.Ranges( + lin_vel_x=(0.0, 1.5), + heading=(-math.pi/4, math.pi/4), + ), + ) + } + + ## + # Events + ## + + events = { + "push_robot": EventTermCfg( + func=mdp.push_by_setting_velocity, + mode="interval", + interval_range_s=(5.0, 10.0), + params={ + "velocity_range": { + "x": (-0.5, 0.5), + "y": (-0.5, 0.5), + }, + }, + ), + "reset_robot_joints": EventTermCfg( + func=mdp.reset_joints_by_offset, + mode="reset", + params={ + "position_range": (-0.01, 0.01), + "velocity_range": (0.0, 0.0), + "asset_cfg": SceneEntityCfg("robot", joint_names=(".*",)), + }, + ), + "base_com": EventTermCfg( + mode="startup", + func=mdp.randomize_field, + domain_randomization=True, + params={ + "asset_cfg": SceneEntityCfg("robot", body_names=("torso_link")), + "operation": "add", + "field": "body_ipos", + "ranges": { + 0: (-0.025, 0.025), + 1: (-0.025, 0.025), + 2: (-0.03, 0.03), + }, + }, + ), + "skate_com": EventTermCfg( + mode="startup", + func=mdp.randomize_field, + domain_randomization=True, + params={ + "asset_cfg": SceneEntityCfg("skateboard", body_names=("skateboard_deck")), + "operation": "add", + "field": "body_ipos", + "ranges": { + 0: (-0.02, 0.02), + 1: (-0.02, 0.02), + 2: (-0.01, 0.01), + }, + }, + ), + "robot_friction": EventTermCfg( + mode="startup", + func=mdp.randomize_field, + domain_randomization=True, + params={ + "asset_cfg": SceneEntityCfg("robot", geom_names=(".*",)), + "operation": "scale", + "field": "geom_friction", + "ranges": (0.3, 1.6), + }, + ), + "board_friction": EventTermCfg( + mode="startup", + func=mdp.randomize_field, + domain_randomization=True, + params={ + "asset_cfg": SceneEntityCfg("skateboard", geom_names=(".*_deck_collision",)), + "operation": "scale", + "field": "geom_friction", + "ranges": (0.8, 2.0), + }, + ), + "foot_friction": EventTermCfg( + mode="startup", + func=mdp.randomize_field, + domain_randomization=True, + params={ + "asset_cfg": SceneEntityCfg("robot", geom_names=(r"^(left|right)_foot[1-7]_collision$")), # Set per-robot. + "operation": "abs", + "field": "geom_friction", + "ranges": (0.3, 1.8), + }, + ), + "wheel_friction": EventTermCfg( + mode="startup", + func=mdp.randomize_field, + domain_randomization=True, + params={ + "asset_cfg": SceneEntityCfg("skateboard", geom_names=(".*_wheel_collision",)), + "operation": "scale", + "field": "geom_friction", + "axes":[2], + "ranges": (0.8, 1.6), + }, + ), + } + + ## + # Rewards + ## + ### push phase rewards + push_rewards = { + "push_skateboard_lin_vel": RewardTermCfg( + func=mdp.push_skateboard_lin_vel, + weight=3.0, + params={"asset_cfg": SceneEntityCfg("skateboard"),"command_name": "skate", "std": math.sqrt(0.25)}, + ), + "push_yaw_align": RewardTermCfg( + func=mdp.push_yaw_align, + weight=1.0, + params={"std": math.sqrt(0.25)}, + ), + "push_air_time": RewardTermCfg( + func=mdp.feet_air_time, + weight=3.0, + params={ + "sensor_name": "left_feet_ground_contact", + "threshold_min": 0.1, + "threshold_max": 0.5, + "command_name": "skate", + "command_threshold": 0.1, + }, + ), + "push_contact_ground_parallel": RewardTermCfg(func=mdp.push_contact_ground_parallel, weight=0.5), + } + ### steer phase rewards + steer_rewards = { + "steer_contact_num": RewardTermCfg( + func=mdp.steer_contact_num, + weight=3.0, + ), + "steer_joint_pos": RewardTermCfg( + func=mdp.steer_joint_pos, + weight=1.5, + params={"std": math.sqrt(0.2)}, + ), + "steer_feet_dis": RewardTermCfg( + func=mdp.steer_feet_dis, + weight=1.0, + params={"std": math.sqrt(0.1)}, + ), + "steer_track_heading": RewardTermCfg( + func=mdp.steer_track_heading, + weight=5.0, + params={"command_name": "skate", "std": math.sqrt(0.02)}, + ), + "steer_tilt_guide": RewardTermCfg( + func=mdp.steer_tilt_guide, + weight=4.0, + params={"command_name": "skate", "std": math.sqrt(0.02)}, + ), + } + ### transition rewards + transition_rewards = { + "transition_body_pos_tracking": RewardTermCfg( + func=mdp.transition_body_pos_tracking, + params={"std": math.sqrt(0.05)}, + weight=10.0, + ), + "transition_body_rot_tracking" : RewardTermCfg( + func=mdp.transition_body_rot_tracking, + params={"std": math.sqrt(0.10)}, + weight=10.0, + ), + "transition_penalty_contact": RewardTermCfg( + func=mdp.transition_penalty_contact, + params={"sensor_name": "left_feet_ground_contact"}, + weight=-0.5, + ), + } + ### regularization rewards + regularization_rewards = { + "reg_wheel_contact_number": RewardTermCfg(func=mdp.reg_wheel_contact_number, weight=0.5), + "dof_pos_limits": RewardTermCfg(func=mdp.joint_pos_limits, weight=-5.0), + "action_rate_l2": RewardTermCfg(func=mdp.action_rate_l2, weight=-0.1), + "action_acc_l2": RewardTermCfg(func=mdp.action_acc_l2, weight=-0.1), + "joint_vel_l2": RewardTermCfg(func=mdp.joint_vel_l2, weight=-1e-3), + "joint_acc_l2": RewardTermCfg(func=mdp.joint_acc_l2, weight=-2.5e-7), + "joint_torques_l2": RewardTermCfg(func=mdp.joint_torques_l2, weight=-1e-6), + "self_collisions": RewardTermCfg(func=mdp.self_collision_cost, weight=-10.0, params={"sensor_name": "robot_collision"}), + "board_flat": RewardTermCfg(func=mdp.board_flat, weight=3.0, params={"std": math.sqrt(0.05)}), + "stand_still": RewardTermCfg(func=mdp.stand_still, weight=1.0, params={"std": math.sqrt(0.1)}), + } + ## + # Terminations + ## + + terminations = { + "time_out": TerminationTermCfg(func=mdp.time_out, time_out=True), + "fell_over": TerminationTermCfg(func=mdp.bad_orientation, params={"limit_angle": math.radians(70.0)}), + "feet_off_board": TerminationTermCfg(func=mdp.bad_feet_off_board), + "illegal_contact": TerminationTermCfg(func=mdp.illegal_contact, params={"sensor_name": "illegal_contact"}), + } + + ## + # Curriculum + ## + + curriculum = { + } + + ## + # Assemble and return + ## + + return G1SkaterManagerBasedRlEnvCfg( + scene=SceneCfg( + terrain=TerrainImporterCfg( + terrain_type="plane", + terrain_generator=None, + ), + num_envs=1, + extent=2.0, + ), + observations=observations, + actions=actions, + commands=commands, + events=events, + terminations=terminations, + curriculum=curriculum, + push_rewards=push_rewards, + steer_rewards=steer_rewards, + transition_rewards=transition_rewards, + regularization_rewards=regularization_rewards, + viewer=ViewerConfig( + origin_type=ViewerConfig.OriginType.ASSET_BODY, + entity_name="robot", + body_name="", # Set per-robot. + distance=4.0, + elevation=-10.0, + azimuth=210.0, + ), + sim=SimulationCfg( + nconmax=35, + njmax=1500, + mujoco=MujocoCfg( + timestep=0.005, + iterations=10, + ls_iterations=20, + ), + ), + decimation=4, + episode_length_s=20.0, + ) diff --git a/src/mjlab_husky/viewer/__init__.py b/src/mjlab_husky/viewer/__init__.py new file mode 100644 index 0000000..cfe9340 --- /dev/null +++ b/src/mjlab_husky/viewer/__init__.py @@ -0,0 +1,4 @@ +from mjlab_husky.viewer.rerun_play_viewer import RerunPlayViewer, RerunPlayViewerCfg + +__all__ = ["RerunPlayViewer", "RerunPlayViewerCfg"] + diff --git a/src/mjlab_husky/viewer/__pycache__/__init__.cpython-312.pyc b/src/mjlab_husky/viewer/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..df26c75 Binary files /dev/null and b/src/mjlab_husky/viewer/__pycache__/__init__.cpython-312.pyc differ diff --git a/src/mjlab_husky/viewer/__pycache__/rerun_native_play_viewer.cpython-312.pyc b/src/mjlab_husky/viewer/__pycache__/rerun_native_play_viewer.cpython-312.pyc new file mode 100644 index 0000000..1f59d03 Binary files /dev/null and b/src/mjlab_husky/viewer/__pycache__/rerun_native_play_viewer.cpython-312.pyc differ diff --git a/src/mjlab_husky/viewer/__pycache__/rerun_play_viewer.cpython-312.pyc b/src/mjlab_husky/viewer/__pycache__/rerun_play_viewer.cpython-312.pyc new file mode 100644 index 0000000..7761b0d Binary files /dev/null and b/src/mjlab_husky/viewer/__pycache__/rerun_play_viewer.cpython-312.pyc differ diff --git a/src/mjlab_husky/viewer/__pycache__/rerun_viser_play_viewer.cpython-312.pyc b/src/mjlab_husky/viewer/__pycache__/rerun_viser_play_viewer.cpython-312.pyc new file mode 100644 index 0000000..3b18016 Binary files /dev/null and b/src/mjlab_husky/viewer/__pycache__/rerun_viser_play_viewer.cpython-312.pyc differ diff --git a/src/mjlab_husky/viewer/rerun_native_play_viewer.py b/src/mjlab_husky/viewer/rerun_native_play_viewer.py new file mode 100644 index 0000000..2c60be9 --- /dev/null +++ b/src/mjlab_husky/viewer/rerun_native_play_viewer.py @@ -0,0 +1,54 @@ +"""MuJoCo 原生 passive 窗口 + Rerun Web：同一套 mjlab 仿真、同一策略步进。""" + +from __future__ import annotations + +import os +from typing import Any + +from mjlab.viewer.base import VerbosityLevel +from mjlab.viewer.native import NativeMujocoViewer + +from mjlab_husky.viewer.rerun_play_viewer import RerunPlayViewer, RerunPlayViewerCfg + + +class RerunNativePlayViewer(NativeMujocoViewer): + """在 `NativeMujocoViewer` 的 `BaseViewer` 主循环上叠加 Rerun 时间轴与相机/qpos 日志。 + + 要求有图形会话（``DISPLAY`` / ``WAYLAND_DISPLAY``），否则无法启动 ``launch_passive``。 + 离屏 ``rgb_array`` 给 Rerun 的相机与 MuJoCo 窗口使用同一份 ``env.unwrapped.sim`` 状态。 + """ + + def __init__( + self, + env: Any, + policy: Any, + rerun_cfg: RerunPlayViewerCfg, + frame_rate: float = 60.0, + verbosity: int = VerbosityLevel.INFO, + ) -> None: + super().__init__(env, policy, frame_rate=frame_rate, verbosity=verbosity) + self._rerun = RerunPlayViewer(env, policy, rerun_cfg) + + def setup(self) -> None: + if not (os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY")): + raise RuntimeError( + "viewer=rerun_native 需要 **MuJoCo 自带的 GLFW 原生窗口**（仓库 `NativeMujocoViewer`），" + "因此进程里必须有 DISPLAY / WAYLAND_DISPLAY。\n" + "纯 SSH / RoboHub 技能容器里若未注入显示环境变量，此处为空就会失败；可选：\n" + " • 只要 Rerun、不要 GLFW 面板： --viewer rerun\n" + " • 虚拟屏幕（常见无桌面服务器）： sudo apt-get install -y xvfb\n" + " xvfb-run -a uv run play ... --viewer rerun_native ...\n" + " • 本机有桌面 / 远程桌面：先 export DISPLAY=:0（或实际编号），再运行。\n" + "当前 DISPLAY/WAYLAND_DISPLAY 均为未设置。" + ) + print( + "[Rerun+Native] 单后端：策略与物理与 `play --viewer native` 相同，" + "并额外把离屏相机与 qpos 写入 Rerun（见终端里的 http 链接）。" + ) + self._rerun.start_rerun_servers(app_id="mjlab_husky_play_rerun_native") + super().setup() + self.reset_environment() + + def sync_env_to_viewer(self) -> None: + super().sync_env_to_viewer() + self._rerun.log_post_step_to_rerun(self._step_count) diff --git a/src/mjlab_husky/viewer/rerun_play_viewer.py b/src/mjlab_husky/viewer/rerun_play_viewer.py new file mode 100644 index 0000000..d264393 --- /dev/null +++ b/src/mjlab_husky/viewer/rerun_play_viewer.py @@ -0,0 +1,370 @@ +from __future__ import annotations + +import os +import re +import shutil +import time +import webbrowser +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from urllib.parse import quote + +import numpy as np +import torch + +import rerun as rr + +from mjlab_husky.lerobot_numpy import batched_vector_to_numpy1d, import_lerobot_dataset_class + + +def _downsample_max_side(img: np.ndarray, max_side: int) -> np.ndarray: + """Cheap decimation so gRPC / web viewer are not flooded with full-HD frames.""" + if max_side <= 0: + return img + h, w = int(img.shape[0]), int(img.shape[1]) + m = max(h, w) + if m <= max_side: + return img + step = max(2, int(np.ceil(m / float(max_side)))) + return np.ascontiguousarray(img[::step, ::step, :]) + + +@dataclass(frozen=True) +class RerunPlayViewerCfg: + web_port: int = 8080 + open_browser: bool = True + grpc_port: int | None = None + connect_host: str | None = None + log_qpos: bool = True + log_qpos_joint_scalars: bool = True + log_qpos_text: bool = False + qpos_text_decimals: int = 3 + log_qpos_tensor: bool = False + qpos_env_idx: int = 0 + log_camera: bool = True + camera_log_stride: int = 2 + qpos_log_stride: int = 4 + camera_max_side: int = 640 + grpc_newest_first: bool = True + lerobot_record: bool = False + lerobot_out_dir: str = "lerobot_data" + lerobot_dataset_name: str = "mjlab_husky_live" + lerobot_steps_per_episode: int = 1000 + lerobot_task: str = "Skater" + lerobot_overwrite: bool = False + lerobot_record_stride: int = 1 + lerobot_save_on_env_reset: bool = False + + +def _viewer_url_with_connect(*, host: str, web_port: int, connect_uri: str) -> str: + return f"http://{host}:{web_port}/?url={quote(connect_uri, safe='')}" + + +def _grpc_listen_port(connect_uri: str) -> int | None: + m = re.search(r"127\.0\.0\.1:(\d+)", connect_uri) + if not m: + return None + return int(m.group(1)) + + +def _grpc_port_from_uri(connect_uri: str) -> int | None: + m = re.search(r":(\d+)/", connect_uri) + if not m: + return None + return int(m.group(1)) + + +def _rewrite_loopback_in_rerun_uri(uri: str, host: str) -> str: + return re.sub(r"127\.0\.0\.1", host.strip(), uri, count=1) + + +class RerunPlayViewer: + def __init__(self, env: Any, policy: Any, cfg: RerunPlayViewerCfg | None = None): + self._env = env + self._policy = policy + self._cfg = cfg or RerunPlayViewerCfg() + self._camera_warned: bool = False + self._lerobot_dataset: Any = None + self._lerobot_steps_in_ep: int = 0 + + def _qpos_to_numpy1d(self, qpos: Any) -> np.ndarray | None: + if qpos is None: + return None + env_i = self._cfg.qpos_env_idx + try: + if isinstance(qpos, torch.Tensor): + t = qpos[env_i].detach().float().cpu().numpy() + elif hasattr(qpos, "detach"): + t = qpos[env_i].detach().float().cpu().numpy() + elif hasattr(qpos, "__getitem__") and hasattr(qpos, "cpu"): + t = qpos[env_i].float().cpu().numpy() + else: + arr = np.asarray(qpos, dtype=np.float64) + if arr.ndim >= 2: + arr = arr[env_i] + t = np.asarray(arr, dtype=np.float64) + return np.asarray(t, dtype=np.float64).reshape(-1) + except Exception: + return None + + def _try_log_mujoco_state(self, step: int): + if not self._cfg.log_qpos: + return + if step % self._cfg.qpos_log_stride != 0: + return + + unwrapped = getattr(self._env, "unwrapped", self._env) + sim = getattr(unwrapped, "sim", None) + data = getattr(sim, "data", None) if sim is not None else None + qpos = getattr(data, "qpos", None) if data is not None else None + + q1d = self._qpos_to_numpy1d(qpos) + if q1d is None or q1d.size == 0: + return + + if self._cfg.log_qpos_tensor: + rr.log("mujoco/qpos/tensor", rr.Tensor(q1d.astype(np.float32))) + + if self._cfg.log_qpos_text: + d = max(0, min(8, int(self._cfg.qpos_text_decimals))) + parts = [f"{float(v):.{d}f}" for v in q1d] + rr.log("mujoco/qpos/values", rr.TextLog(" ".join(parts))) + + if self._cfg.log_qpos_joint_scalars: + for i, v in enumerate(q1d): + rr.log(f"mujoco/qpos/j_{i:03d}", rr.Scalars([float(v)])) + + def _try_log_camera(self, step: int): + if not self._cfg.log_camera: + return + if step % self._cfg.camera_log_stride != 0: + return + unwrapped = getattr(self._env, "unwrapped", self._env) + if getattr(unwrapped, "render_mode", None) != "rgb_array": + if not self._camera_warned: + print( + "[Rerun] No rgb_array render_mode on env — camera disabled. " + "Use `play` with Rerun viewer so OffscreenRenderer is enabled." + ) + self._camera_warned = True + return + try: + frame = unwrapped.render() + except Exception as e: + if not self._camera_warned: + print(f"[Rerun] render() failed ({e}); camera disabled.") + self._camera_warned = True + return + if frame is None: + return + img = np.ascontiguousarray(frame) + if img.dtype != np.uint8: + img = (np.clip(img, 0.0, 1.0) * 255.0).astype(np.uint8) + img = _downsample_max_side(img, self._cfg.camera_max_side) + rr.log("camera/robot", rr.Image(img)) + + def start_rerun_servers(self, *, app_id: str = "mjlab_husky_play_rerun") -> None: + """Initialize Rerun + gRPC + web viewer (no sim loop). Used by ``run()`` and ``RerunNativePlayViewer``.""" + rr.init(app_id, spawn=False) + + sg_kw: dict[str, Any] = {"newest_first": self._cfg.grpc_newest_first} + if self._cfg.grpc_port is not None: + sg_kw["grpc_port"] = self._cfg.grpc_port + server_uri = rr.serve_grpc(**sg_kw) + + connect_uri = server_uri + viewer_http_host = "127.0.0.1" + if self._cfg.connect_host: + connect_uri = _rewrite_loopback_in_rerun_uri(server_uri, self._cfg.connect_host) + viewer_http_host = self._cfg.connect_host.strip() + + local_viewer = _viewer_url_with_connect( + host=viewer_http_host, + web_port=self._cfg.web_port, + connect_uri=connect_uri, + ) + grpc_port = _grpc_listen_port(server_uri) or _grpc_port_from_uri(connect_uri) + print("[Rerun] gRPC data source (serve_grpc):", server_uri) + if connect_uri != server_uri: + print("[Rerun] gRPC URL passed to web viewer (connect_to):", connect_uri) + print("[Rerun] Open this URL in the browser (required to see live data):", local_viewer) + print( + "[Rerun] 重要：`?url=` 里的「+」必须是 **%2B**（见上面链接）。" + "若在地址栏手写成了 `rerun+http://...`，`+` 会被当成空格并报 Failed to parse URL。" + ) + print( + "[Rerun] 请复制终端里以 `http://` 开头的整行到浏览器；勿点单独的「rerun…」伪链接。" + ) + if self._cfg.log_camera: + print( + "[Rerun] Camera -> `camera/robot` " + f"(stride={self._cfg.camera_log_stride}, max_side={self._cfg.camera_max_side}); " + f"qpos stride={self._cfg.qpos_log_stride}." + ) + if grpc_port is not None: + print( + "[Rerun] 若浏览器与运行 play 的机器不是同一台：在本机执行 SSH 转发（Web+gRPC 两个端口），" + "否则会 Failed to fetch。", + ) + print( + "[Rerun] SSH example:", + f"ssh -N -L {self._cfg.web_port}:127.0.0.1:{self._cfg.web_port} " + f"-L {grpc_port}:127.0.0.1:{grpc_port} user@remote", + ) + if os.environ.get("SSH_CONNECTION") and not self._cfg.connect_host: + print( + "[Rerun] 当前在 SSH 会话中：上面的 ssh -L 要在「你自己的电脑」执行，不是在这台云主机里执行。", + ) + + if self._cfg.open_browser: + try: + webbrowser.open(local_viewer) + except Exception as e: + print(f"[Rerun] webbrowser.open failed ({e}); 请手动复制上面的 http:// 链接。") + + rr.serve_web_viewer( + web_port=self._cfg.web_port, + open_browser=False, + connect_to=connect_uri, + ) + + def log_post_step_to_rerun(self, step: int) -> None: + """After ``env.step``, stream qpos / offscreen camera to Rerun (same timeline as standalone ``run``).""" + rr.set_time("step", sequence=step) + self._try_log_mujoco_state(step) + self._try_log_camera(step) + + def run(self): + self.start_rerun_servers() + + reset_out = self._env.reset() + obs = reset_out[0] if isinstance(reset_out, tuple) and len(reset_out) == 2 else reset_out + step = 0 + + if self._cfg.lerobot_record: + LeRobotDataset = import_lerobot_dataset_class() + + unwrapped = getattr(self._env, "unwrapped", self._env) + sim_data = getattr(getattr(unwrapped, "sim", None), "data", None) + if sim_data is None or getattr(sim_data, "qpos", None) is None: + raise RuntimeError("LeRobot recording needs env.unwrapped.sim.data.qpos") + q0 = self._qpos_to_numpy1d(sim_data.qpos) + if q0 is None: + raise RuntimeError("Could not read qpos for LeRobot schema") + nq = int(q0.shape[0]) + na = int(np.prod(np.asarray(unwrapped.action_space.shape, dtype=np.int64))) + step_dt = float(getattr(unwrapped, "step_dt", 0.02)) + fps = max(1, int(round(1.0 / step_dt))) + ds_root = (Path.cwd() / self._cfg.lerobot_out_dir / self._cfg.lerobot_dataset_name).resolve() + if ds_root.exists(): + if not self._cfg.lerobot_overwrite: + raise FileExistsError(f"LeRobot dataset exists: {ds_root} (use --lerobot-overwrite)") + shutil.rmtree(ds_root) + ds_root.parent.mkdir(parents=True, exist_ok=True) + self._lerobot_dataset = LeRobotDataset.create( + repo_id=self._cfg.lerobot_dataset_name, + fps=fps, + features={ + "observation.state": {"dtype": "float32", "shape": (nq,), "names": None}, + "action": {"dtype": "float32", "shape": (na,), "names": None}, + }, + root=str(ds_root), + robot_type="mjlab_husky", + use_videos=False, + ) + self._lerobot_steps_in_ep = 0 + print(f"[LeRobot] Recording enabled -> {ds_root} (fps={fps}, nq={nq}, na={na})") + + try: + while True: + rr.set_time("step", sequence=step) + + with np.errstate(all="ignore"): + action = self._policy(obs) + + if self._lerobot_dataset is not None and step % max(1, self._cfg.lerobot_record_stride) == 0: + u0 = getattr(self._env, "unwrapped", self._env) + sd0 = getattr(getattr(u0, "sim", None), "data", None) + qpos_pre = getattr(sd0, "qpos", None) if sd0 is not None else None + try: + q_rec = ( + batched_vector_to_numpy1d(qpos_pre, env_i=self._cfg.qpos_env_idx) + if qpos_pre is not None + else None + ) + a_rec = batched_vector_to_numpy1d(action, env_i=self._cfg.qpos_env_idx) + except Exception: + q_rec, a_rec = None, None + if q_rec is not None and a_rec is not None: + self._lerobot_dataset.add_frame( + { + "task": self._cfg.lerobot_task, + "observation.state": q_rec, + "action": a_rec, + } + ) + self._lerobot_steps_in_ep += 1 + if self._lerobot_steps_in_ep >= int(self._cfg.lerobot_steps_per_episode): + self._lerobot_dataset.save_episode() + self._lerobot_steps_in_ep = 0 + print("[LeRobot] Saved episode shard (save_episode).") + + step_out = self._env.step(action) + if isinstance(step_out, tuple) and len(step_out) == 5: + obs, rew, terminated, truncated, info = step_out + try: + done = np.asarray(terminated) | np.asarray(truncated) + except Exception: + done = bool(terminated) or bool(truncated) + else: + obs, rew, done, info = step_out + + try: + rew_mean = float(np.asarray(rew).mean()) + rr.log("env/reward_mean", rr.Scalars([rew_mean])) + except Exception: + pass + + try: + done_frac = float(np.asarray(done, dtype=np.float32).mean()) + rr.log("env/done_frac", rr.Scalars([done_frac])) + except Exception: + pass + + self._try_log_mujoco_state(step) + self._try_log_camera(step) + + try: + if bool(np.any(done)): + if ( + self._lerobot_dataset is not None + and self._cfg.lerobot_save_on_env_reset + and self._lerobot_steps_in_ep > 0 + ): + self._lerobot_dataset.save_episode() + self._lerobot_steps_in_ep = 0 + print("[LeRobot] save_episode() on env reset (save_on_env_reset).") + reset_out = self._env.reset() + obs = ( + reset_out[0] + if isinstance(reset_out, tuple) and len(reset_out) == 2 + else reset_out + ) + except Exception: + pass + + step += 1 + time.sleep(0.0) + finally: + if self._lerobot_dataset is not None: + try: + if self._lerobot_steps_in_ep > 0: + self._lerobot_dataset.save_episode() + except Exception as e: + print(f"[LeRobot] Warning: final save_episode failed: {e}") + try: + self._lerobot_dataset.finalize() + except Exception as e: + print(f"[LeRobot] Warning: finalize failed: {e}") + print("[LeRobot] finalize() done.") diff --git a/src/mjlab_husky/viewer/rerun_viser_play_viewer.py b/src/mjlab_husky/viewer/rerun_viser_play_viewer.py new file mode 100644 index 0000000..3e0c442 --- /dev/null +++ b/src/mjlab_husky/viewer/rerun_viser_play_viewer.py @@ -0,0 +1,65 @@ +"""Viser（浏览器里的 mjlab 三维面板）+ Rerun：单套仿真，两个服务各占不同 TCP 端口。""" + +from __future__ import annotations + +from typing import Any + +import viser +from typing_extensions import override + +from mjlab.viewer.base import VerbosityLevel +from mjlab.viewer.viser import ViserPlayViewer + +from mjlab_husky.viewer.rerun_play_viewer import RerunPlayViewer, RerunPlayViewerCfg + + +class RerunViserPlayViewer(ViserPlayViewer): + """与 ``ViserPlayViewer`` 相同的主循环，额外启动 Rerun Web+gRPC，并在每步写入时间轴。 + + - **Viser** 占用 ``--viser-port``（默认 `8081`，避免与 Rerun Web 默认 `8080` 冲突）。 + - **Rerun** 占用 ``--rerun-web-port`` 与 ``--rerun-grpc-port``（与 ``--viewer rerun`` 一致）。 + + GLFW ``NativeMujocoViewer`` 不监听端口；若要在浏览器里并排嵌入「三维 + Rerun」，请用本 viewer 而非 ``rerun_native``。 + """ + + def __init__( + self, + env: Any, + policy: Any, + rerun_cfg: RerunPlayViewerCfg, + viser_port: int = 8081, + frame_rate: float = 60.0, + verbosity: int = VerbosityLevel.SILENT, + ) -> None: + super().__init__(env, policy, frame_rate=frame_rate, verbosity=verbosity) + self._rerun = RerunPlayViewer(env, policy, rerun_cfg) + self._viser_port = int(viser_port) + self._url_host = (rerun_cfg.connect_host or "127.0.0.1").strip() + + @override + def setup(self) -> None: + self._rerun.start_rerun_servers(app_id="mjlab_husky_play_rerun_viser") + + _real = viser.ViserServer + port = self._viser_port + + def _force_port_ViserServer(*args: Any, **kwargs: Any): + kwargs = dict(kwargs) + kwargs["port"] = port + return _real(*args, **kwargs) + + viser.ViserServer = _force_port_ViserServer # type: ignore[misc, assignment] + try: + super().setup() + finally: + viser.ViserServer = _real + + print( + f"[Rerun+Viser] 单后端：Rerun 使用 --rerun-web-port / --rerun-grpc-port；" + f"Viser 三维面板: http://{self._url_host}:{port}/" + ) + + @override + def sync_env_to_viewer(self) -> None: + super().sync_env_to_viewer() + self._rerun.log_post_step_to_rerun(self._step_count) diff --git a/test_scene/mjlab_scene.xml b/test_scene/mjlab_scene.xml new file mode 100644 index 0000000..e15d0db --- /dev/null +++ b/test_scene/mjlab_scene.xml @@ -0,0 +1,383 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test_scene/sim.py b/test_scene/sim.py new file mode 100644 index 0000000..7017eba --- /dev/null +++ b/test_scene/sim.py @@ -0,0 +1,374 @@ +import argparse +import math +import time +import numpy as np +import mujoco +import torch +from rich import print +from collections import deque +import mujoco.viewer as mjv +from tqdm import tqdm +import os + +try: + import onnxruntime as ort +except ImportError: + ort = None + +class OnnxPolicyWrapper: + """Minimal wrapper so ONNXRuntime policies mimic TorchScript call signature.""" + + def __init__(self, session, input_name, output_index=0): + self.session = session + self.input_name = input_name + self.output_index = output_index + + def __call__(self, obs_tensor: torch.Tensor) -> torch.Tensor: + if isinstance(obs_tensor, torch.Tensor): + obs_np = obs_tensor.detach().cpu().numpy() + else: + obs_np = np.asarray(obs_tensor, dtype=np.float32) + outputs = self.session.run(None, {self.input_name: obs_np}) + result = outputs[self.output_index] + if not isinstance(result, np.ndarray): + result = np.asarray(result, dtype=np.float32) + return torch.from_numpy(result.astype(np.float32)) + + +def load_onnx_policy(policy_path: str, device: str) -> OnnxPolicyWrapper: + if ort is None: + raise ImportError("onnxruntime is required for ONNX policy inference but is not installed.") + providers = [] + available = ort.get_available_providers() + if device.startswith('cuda'): + if 'CUDAExecutionProvider' in available: + providers.append('CUDAExecutionProvider') + else: + print("CUDAExecutionProvider not available in onnxruntime; falling back to CPUExecutionProvider.") + providers.append('CPUExecutionProvider') + session = ort.InferenceSession(policy_path, providers=providers) + input_name = session.get_inputs()[0].name + print(f"ONNX policy loaded from {policy_path} using providers: {session.get_providers()}") + return OnnxPolicyWrapper(session, input_name) + +from pynput import keyboard +import threading + +reset_flag = False +pause_flag = False +V_MIN, V_MAX = 0.0, 1.5 +H_MIN, H_MAX = -math.pi / 4, math.pi / 4 +v = 1.0 +h = 0.0 + + +def wrap_to_pi(x): + return (x + math.pi) % (2.0 * math.pi) - math.pi + + +def on_press(key): + global v, h, reset_flag, pause_flag + try: + if key == keyboard.Key.up: + v = round(min(v + 0.1, V_MAX), 1) + print("v =", v, "h =", round(h, 3), "(rad)") + elif key == keyboard.Key.down: + v = round(max(v - 0.1, V_MIN), 1) + print("v =", v, "h =", round(h, 3), "(rad)") + elif key == keyboard.Key.left: + h = round(max(h + 0.1, H_MIN), 2) + print("v =", v, "h =", round(h, 3), "(rad)") + elif key == keyboard.Key.right: + h = round(min(h - 0.1, H_MAX), 2) + print("v =", v, "h =", round(h, 3), "(rad)") + elif key == keyboard.Key.enter: + reset_flag = True + print("Reset flag set! Simulation will reset...") + elif key == keyboard.Key.space: + pause_flag = not pause_flag + if pause_flag: + print("Simulation PAUSED. Press SPACE to resume.") + else: + print("Simulation RESUMED.") + elif hasattr(key, "char") and key.char == "5": + v = 0.0 + h = 0.0 + print("Commands reset: v = 0.0, h = 0.0") + except AttributeError: + pass + +def start_listener(): + with keyboard.Listener(on_press=on_press) as listener: + listener.join() + +listener_thread = threading.Thread(target=start_listener) +listener_thread.daemon = True +listener_thread.start() + +def get_gravity_orientation(quaternion): + qw = quaternion[0] + qx = quaternion[1] + qy = quaternion[2] + qz = quaternion[3] + + gravity_orientation = np.zeros(3) + + gravity_orientation[0] = 2 * (-qz * qx + qw * qy) + gravity_orientation[1] = -2 * (qz * qy + qw * qx) + gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz) + + return gravity_orientation + +def quat_apply_np(quat, vec): + quat = np.asarray(quat) + vec = np.asarray(vec) + orig_shape = vec.shape + + q = quat.reshape(-1, 4) + v = vec.reshape(-1, 3) + + w = q[:, 0] + qvec = q[:, 1:4] + + t = 2 * np.cross(qvec, v) + v_rot = v + (w[:, None] * t) + np.cross(qvec, t) + v_rot = v_rot.reshape(orig_shape) + return v_rot + +reindex_list = [15, 16, 17, 18, 19, 20, 21, 22, 0, 2, 6, 8, 12, 1, 3, 7, 9, 13, 14, 4, 5, 10, 11] + +class RealTimePolicyController: + def __init__(self, + xml_file, + policy_path, + device='cuda', + policy_frequency=50, + ): + + self.device = device + self.policy = load_onnx_policy(policy_path, device) + + # Create MuJoCo sim + self.model = mujoco.MjModel.from_xml_path(xml_file) + self.model.opt.timestep = 0.005 + self.model.opt.iterations = 10 + self.model.opt.ls_iterations = 20 + self.model.opt.ccd_iterations = 50 + + self.data = mujoco.MjData(self.model) + + self.viewer = mjv.launch_passive(self.model, self.data, show_left_ui=False, show_right_ui=False) + self.viewer.cam.distance = 4.0 + self.viewer.cam.azimuth = 210.0 + self.viewer.cam.elevation = -10.0 + self.num_actions = 23 + self.sim_duration = 30.0 + self.sim_dt = 0.005 + self.cycle_time = 6 + self.step_dt = 1 / policy_frequency + self.sim_decimation = int(1 / (policy_frequency * self.sim_dt)) + + print(f"sim_decimation: {self.sim_decimation}") + + self.last_action = np.zeros(self.num_actions, dtype=np.float32) + + self.robot_default_dof_pos = np.array([ + 0.0, 0.0, 0.0, 0.23, -0.20, 0.0, + -0.7, 0.0, 0.0, 1.17, -0.45, 0.0, + 0.0, 0.0, 0.0, + -0.03, 0.45, -0.21, 1.32, + -0.7, -0.845, 0.83, 1.19 + ]) + + self.mujoco_default_dof_pos = np.concatenate([ + np.array([-0.03, 0.1, 0.78]), + np.array([1, 0, 0, 0]), + self.robot_default_dof_pos, + np.array([0, 0, 0.10]), + np.array([1, 0, 0, 0]), + np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + ]) + + self.action_scale = np.array([ + 0.5475, 0.3507, 0.5475, 0.3507, 0.4386, 0.4386, + 0.5475, 0.3507, 0.5475, 0.3507, 0.4386, 0.4386, + 0.5475, 0.4386, 0.4386, + 0.4386, 0.4386, 0.4386, 0.4386, + 0.4386, 0.4386, 0.4386, 0.4386, + ]) + + self.n_obs_single = 3 + 3 + 3 + 3*23 + 1 + self.history_len = 5 + self.total_obs_size = self.n_obs_single * (self.history_len) + + self.obs_block_dims = [2, 1, 3, 3, 23, 23, 23, 1] + self.obs_block_starts = np.cumsum([0] + self.obs_block_dims[:-1]) + + self.proprio_history_buf = deque(maxlen=self.history_len) + for _ in range(self.history_len): + self.proprio_history_buf.append(np.zeros(self.n_obs_single, dtype=np.float32)) + + def reset_sim(self): + """Reset simulation to initial state""" + mujoco.mj_resetData(self.model, self.data) + mujoco.mj_forward(self.model, self.data) + + def reset(self, init_pos): + """Reset robot to initial position""" + self.data.qpos[:] = init_pos + self.data.qvel[:] = 0 + self.data.ctrl[:-7] = self.robot_default_dof_pos[reindex_list] + mujoco.mj_forward(self.model, self.data) + + def extract_data(self): + n_robot_dof = self.num_actions + + robot_quat = self.data.qpos[3:7] + robot_dof_pos = self.data.qpos[7:7+n_robot_dof] + robot_ang_vel = self.data.qvel[3:6] + robot_dof_vel = self.data.qvel[6:6+n_robot_dof] + + return robot_dof_pos, robot_dof_vel, robot_quat, robot_ang_vel + + def run(self): + """Main simulation loop""" + global reset_flag, pause_flag, v, h + print("Starting Skater simulation...") + + self.reset_sim() + self.reset(self.mujoco_default_dof_pos) + + steps = int(self.sim_duration / self.sim_dt) + pbar = tqdm(range(steps), desc="Simulating Skater...") + + phase_counter = 0 + + try: + for i in pbar: + if not self.viewer.is_running(): + print("Viewer closed, stopping simulation.") + break + if reset_flag: + self.reset_sim() + self.reset(self.mujoco_default_dof_pos) + reset_flag = False + phase_counter = 0 + print("Simulation RESET!") + if pause_flag: + time.sleep(0.01) + continue + t_start = time.time() + + phase_counter += 1 + + phase = ((phase_counter * self.step_dt / self.cycle_time)) % 1.0 + phase = torch.tensor(phase) + phase = torch.clip(phase, 0.0, 1.0) + + robot_dof_pos, robot_dof_vel, robot_quat, robot_ang_vel = self.extract_data() + + gravity_orientation = get_gravity_orientation(robot_quat) + + sensor_id = self.model.sensor("robot/imu_ang_vel").id + sensor_adr = self.model.sensor_adr[sensor_id] + sensor_dim = self.model.sensor_dim[sensor_id] + sensor_ang_vel = self.data.sensordata[sensor_adr : sensor_adr + sensor_dim] + + forward_w = quat_apply_np(robot_quat, np.array([1, 0, 0])) + heading = np.array([np.arctan2(forward_w[1], forward_w[0])]) + + obs_proprio = np.concatenate([ + np.array([v, h], dtype=np.float32) * [2.0, 1.0], + heading * 1.0 / math.pi, + sensor_ang_vel * 0.25, + gravity_orientation, + (robot_dof_pos - self.robot_default_dof_pos), + robot_dof_vel * 0.05, + self.last_action, + np.array([phase], dtype=np.float32), + ]) + + self.proprio_history_buf.append(obs_proprio) + history_array = np.array(self.proprio_history_buf) + + obs_buf_parts = [] + for i, (start, dim) in enumerate(zip(self.obs_block_starts, self.obs_block_dims)): + obs_block = history_array[:, start:start+dim] + obs_buf_parts.append(obs_block.flatten()) + + obs_buf = np.concatenate(obs_buf_parts) + + obs_tensor = torch.from_numpy(obs_buf).float().unsqueeze(0).to(self.device) + with torch.no_grad(): + raw_action = self.policy(obs_tensor).cpu().numpy().squeeze() + + self.last_action = raw_action + scaled_actions = raw_action * self.action_scale + + pd_target_robot = (scaled_actions + self.robot_default_dof_pos) + + viewer_closed = False + for _ in range(self.sim_decimation): + if not self.viewer.is_running(): + viewer_closed = True + break + self.data.ctrl[:-7] = pd_target_robot[reindex_list] + mujoco.mj_step(self.model, self.data) + pelvis_pos = self.data.xpos[self.model.body("robot/pelvis").id] + self.viewer.cam.lookat = pelvis_pos + self.viewer.sync() + if viewer_closed: + break + + dt = self.model.opt.timestep * self.sim_decimation + sleep = dt - (time.time() - t_start) + if sleep > 0: + time.sleep(sleep) + + except Exception as e: + print(f"Error in run: {e}") + import traceback + traceback.print_exc() + finally: + + if self.viewer: + self.viewer.close() + print("Simulation finished.") + + +def main(): + parser = argparse.ArgumentParser(description='Run skater policy in simulation') + parser.add_argument('--xml', type=str, default='mjlab_scene.xml', + help='Path to MuJoCo XML file') + parser.add_argument('--policy', type=str, required=True, + help='Path to skater ONNX policy file') + parser.add_argument('--device', type=str, + default='cuda', + help='Device to run policy on (cuda/cpu)') + parser.add_argument("--policy_frequency", help="Policy frequency", default=50, type=int) + args = parser.parse_args() + + if not os.path.exists(args.policy): + print(f"Error: Policy file {args.policy} does not exist") + return + + if not os.path.exists(args.xml): + print(f"Error: XML file {args.xml} does not exist") + return + + print(f"Starting skater simulation controller...") + print(f" XML file: {args.xml}") + print(f" Policy file: {args.policy}") + print(f" Device: {args.device}") + + controller = RealTimePolicyController( + xml_file=args.xml, + policy_path=args.policy, + device=args.device, + policy_frequency=args.policy_frequency, + ) + controller.run() + + +if __name__ == "__main__": + main() diff --git a/test_scene/sim.sh b/test_scene/sim.sh new file mode 100644 index 0000000..dd43481 --- /dev/null +++ b/test_scene/sim.sh @@ -0,0 +1,14 @@ +SCRIPT_DIR=$(dirname $(realpath $0)) + +if [ -z "$1" ]; then + echo "Usage: $0 path_to_policy.onnx" + exit 1 +fi + +ckpt_path=$1 + +uv run python test_scene/sim.py \ + --xml ${SCRIPT_DIR}/mjlab_scene.xml \ + --policy ${ckpt_path} \ + --device cuda \ + --policy_frequency 50 diff --git a/uv.lock b/uv.lock new file mode 100644 index 0000000..2f109c9 --- /dev/null +++ b/uv.lock @@ -0,0 +1,2643 @@ +version = 1 +revision = 3 +requires-python = ">=3.12, <3.14" +resolution-markers = [ + "python_full_version >= '3.13' and platform_machine != 's390x' and sys_platform != 'darwin'", + "python_full_version < '3.13' and platform_machine != 's390x' and sys_platform != 'darwin'", + "python_full_version >= '3.13' and platform_machine == 's390x' and sys_platform != 'darwin'", + "python_full_version < '3.13' and platform_machine == 's390x' and sys_platform != 'darwin'", + "python_full_version >= '3.13' and platform_machine != 's390x' and sys_platform == 'darwin'", + "python_full_version < '3.13' and platform_machine != 's390x' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and platform_machine == 's390x' and sys_platform == 'darwin'", + "python_full_version < '3.13' and platform_machine == 's390x' and sys_platform == 'darwin'", +] + +[[package]] +name = "absl-py" +version = "2.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/64/c7/8de93764ad66968d19329a7e0c147a2bb3c7054c554d4a119111b8f9440f/absl_py-2.4.0.tar.gz", hash = "sha256:8c6af82722b35cf71e0f4d1d47dcaebfff286e27110a99fc359349b247dfb5d4", size = 116543, upload-time = "2026-01-28T10:17:05.322Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/a6/907a406bb7d359e6a63f99c313846d9eec4f7e6f7437809e03aa00fa3074/absl_py-2.4.0-py3-none-any.whl", hash = "sha256:88476fd881ca8aab94ffa78b7b6c632a782ab3ba1cd19c9bd423abc4fb4cd28d", size = 135750, upload-time = "2026-01-28T10:17:04.19Z" }, +] + +[[package]] +name = "annotated-types" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ee/67/531ea369ba64dcff5ec9c3402f9f51bf748cec26dde048a2f973a4eea7f5/annotated_types-0.7.0.tar.gz", hash = "sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89", size = 16081, upload-time = "2024-05-20T21:33:25.928Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, +] + +[[package]] +name = "anyio" +version = "4.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "idna" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/19/14/2c5dd9f512b66549ae92767a9c7b330ae88e1932ca57876909410251fe13/anyio-4.13.0.tar.gz", hash = "sha256:334b70e641fd2221c1505b3890c69882fe4a2df910cba14d97019b90b24439dc", size = 231622, upload-time = "2026-03-24T12:59:09.671Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/42/e921fccf5015463e32a3cf6ee7f980a6ed0f395ceeaa45060b61d86486c2/anyio-4.13.0-py3-none-any.whl", hash = "sha256:08b310f9e24a9594186fd75b4f73f4a4152069e3853f1ed8bfbf58369f4ad708", size = 114353, upload-time = "2026-03-24T12:59:08.246Z" }, +] + +[[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 = "bcrypt" +version = "5.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d4/36/3329e2518d70ad8e2e5817d5a4cac6bba05a47767ec416c7d020a965f408/bcrypt-5.0.0.tar.gz", hash = "sha256:f748f7c2d6fd375cc93d3fba7ef4a9e3a092421b8dbf34d8d4dc06be9492dfdd", size = 25386, upload-time = "2025-09-25T19:50:47.829Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/13/85/3e65e01985fddf25b64ca67275bb5bdb4040bd1a53b66d355c6c37c8a680/bcrypt-5.0.0-cp313-cp313t-macosx_10_12_universal2.whl", hash = "sha256:f3c08197f3039bec79cee59a606d62b96b16669cff3949f21e74796b6e3cd2be", size = 481806, upload-time = "2025-09-25T19:49:05.102Z" }, + { url = "https://files.pythonhosted.org/packages/44/dc/01eb79f12b177017a726cbf78330eb0eb442fae0e7b3dfd84ea2849552f3/bcrypt-5.0.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:200af71bc25f22006f4069060c88ed36f8aa4ff7f53e67ff04d2ab3f1e79a5b2", size = 268626, upload-time = "2025-09-25T19:49:06.723Z" }, + { url = "https://files.pythonhosted.org/packages/8c/cf/e82388ad5959c40d6afd94fb4743cc077129d45b952d46bdc3180310e2df/bcrypt-5.0.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:baade0a5657654c2984468efb7d6c110db87ea63ef5a4b54732e7e337253e44f", size = 271853, upload-time = "2025-09-25T19:49:08.028Z" }, + { url = "https://files.pythonhosted.org/packages/ec/86/7134b9dae7cf0efa85671651341f6afa695857fae172615e960fb6a466fa/bcrypt-5.0.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:c58b56cdfb03202b3bcc9fd8daee8e8e9b6d7e3163aa97c631dfcfcc24d36c86", size = 269793, upload-time = "2025-09-25T19:49:09.727Z" }, + { url = "https://files.pythonhosted.org/packages/cc/82/6296688ac1b9e503d034e7d0614d56e80c5d1a08402ff856a4549cb59207/bcrypt-5.0.0-cp313-cp313t-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:4bfd2a34de661f34d0bda43c3e4e79df586e4716ef401fe31ea39d69d581ef23", size = 289930, upload-time = "2025-09-25T19:49:11.204Z" }, + { url = "https://files.pythonhosted.org/packages/d1/18/884a44aa47f2a3b88dd09bc05a1e40b57878ecd111d17e5bba6f09f8bb77/bcrypt-5.0.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:ed2e1365e31fc73f1825fa830f1c8f8917ca1b3ca6185773b349c20fd606cec2", size = 272194, upload-time = "2025-09-25T19:49:12.524Z" }, + { url = "https://files.pythonhosted.org/packages/0e/8f/371a3ab33c6982070b674f1788e05b656cfbf5685894acbfef0c65483a59/bcrypt-5.0.0-cp313-cp313t-manylinux_2_34_aarch64.whl", hash = "sha256:83e787d7a84dbbfba6f250dd7a5efd689e935f03dd83b0f919d39349e1f23f83", size = 269381, upload-time = "2025-09-25T19:49:14.308Z" }, + { url = "https://files.pythonhosted.org/packages/b1/34/7e4e6abb7a8778db6422e88b1f06eb07c47682313997ee8a8f9352e5a6f1/bcrypt-5.0.0-cp313-cp313t-manylinux_2_34_x86_64.whl", hash = "sha256:137c5156524328a24b9fac1cb5db0ba618bc97d11970b39184c1d87dc4bf1746", size = 271750, upload-time = "2025-09-25T19:49:15.584Z" }, + { url = "https://files.pythonhosted.org/packages/c0/1b/54f416be2499bd72123c70d98d36c6cd61a4e33d9b89562c22481c81bb30/bcrypt-5.0.0-cp313-cp313t-musllinux_1_1_aarch64.whl", hash = "sha256:38cac74101777a6a7d3b3e3cfefa57089b5ada650dce2baf0cbdd9d65db22a9e", size = 303757, upload-time = "2025-09-25T19:49:17.244Z" }, + { url = "https://files.pythonhosted.org/packages/13/62/062c24c7bcf9d2826a1a843d0d605c65a755bc98002923d01fd61270705a/bcrypt-5.0.0-cp313-cp313t-musllinux_1_1_x86_64.whl", hash = "sha256:d8d65b564ec849643d9f7ea05c6d9f0cd7ca23bdd4ac0c2dbef1104ab504543d", size = 306740, upload-time = "2025-09-25T19:49:18.693Z" }, + { url = "https://files.pythonhosted.org/packages/d5/c8/1fdbfc8c0f20875b6b4020f3c7dc447b8de60aa0be5faaf009d24242aec9/bcrypt-5.0.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:741449132f64b3524e95cd30e5cd3343006ce146088f074f31ab26b94e6c75ba", size = 334197, upload-time = "2025-09-25T19:49:20.523Z" }, + { url = "https://files.pythonhosted.org/packages/a6/c1/8b84545382d75bef226fbc6588af0f7b7d095f7cd6a670b42a86243183cd/bcrypt-5.0.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:212139484ab3207b1f0c00633d3be92fef3c5f0af17cad155679d03ff2ee1e41", size = 352974, upload-time = "2025-09-25T19:49:22.254Z" }, + { url = "https://files.pythonhosted.org/packages/10/a6/ffb49d4254ed085e62e3e5dd05982b4393e32fe1e49bb1130186617c29cd/bcrypt-5.0.0-cp313-cp313t-win32.whl", hash = "sha256:9d52ed507c2488eddd6a95bccee4e808d3234fa78dd370e24bac65a21212b861", size = 148498, upload-time = "2025-09-25T19:49:24.134Z" }, + { url = "https://files.pythonhosted.org/packages/48/a9/259559edc85258b6d5fc5471a62a3299a6aa37a6611a169756bf4689323c/bcrypt-5.0.0-cp313-cp313t-win_amd64.whl", hash = "sha256:f6984a24db30548fd39a44360532898c33528b74aedf81c26cf29c51ee47057e", size = 145853, upload-time = "2025-09-25T19:49:25.702Z" }, + { url = "https://files.pythonhosted.org/packages/2d/df/9714173403c7e8b245acf8e4be8876aac64a209d1b392af457c79e60492e/bcrypt-5.0.0-cp313-cp313t-win_arm64.whl", hash = "sha256:9fffdb387abe6aa775af36ef16f55e318dcda4194ddbf82007a6f21da29de8f5", size = 139626, upload-time = "2025-09-25T19:49:26.928Z" }, + { url = "https://files.pythonhosted.org/packages/84/29/6237f151fbfe295fe3e074ecc6d44228faa1e842a81f6d34a02937ee1736/bcrypt-5.0.0-cp38-abi3-macosx_10_12_universal2.whl", hash = "sha256:fc746432b951e92b58317af8e0ca746efe93e66555f1b40888865ef5bf56446b", size = 494553, upload-time = "2025-09-25T19:49:49.006Z" }, + { url = "https://files.pythonhosted.org/packages/45/b6/4c1205dde5e464ea3bd88e8742e19f899c16fa8916fb8510a851fae985b5/bcrypt-5.0.0-cp38-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c2388ca94ffee269b6038d48747f4ce8df0ffbea43f31abfa18ac72f0218effb", size = 275009, upload-time = "2025-09-25T19:49:50.581Z" }, + { url = "https://files.pythonhosted.org/packages/3b/71/427945e6ead72ccffe77894b2655b695ccf14ae1866cd977e185d606dd2f/bcrypt-5.0.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:560ddb6ec730386e7b3b26b8b4c88197aaed924430e7b74666a586ac997249ef", size = 278029, upload-time = "2025-09-25T19:49:52.533Z" }, + { url = "https://files.pythonhosted.org/packages/17/72/c344825e3b83c5389a369c8a8e58ffe1480b8a699f46c127c34580c4666b/bcrypt-5.0.0-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:d79e5c65dcc9af213594d6f7f1fa2c98ad3fc10431e7aa53c176b441943efbdd", size = 275907, upload-time = "2025-09-25T19:49:54.709Z" }, + { url = "https://files.pythonhosted.org/packages/0b/7e/d4e47d2df1641a36d1212e5c0514f5291e1a956a7749f1e595c07a972038/bcrypt-5.0.0-cp38-abi3-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:2b732e7d388fa22d48920baa267ba5d97cca38070b69c0e2d37087b381c681fd", size = 296500, upload-time = "2025-09-25T19:49:56.013Z" }, + { url = "https://files.pythonhosted.org/packages/0f/c3/0ae57a68be2039287ec28bc463b82e4b8dc23f9d12c0be331f4782e19108/bcrypt-5.0.0-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:0c8e093ea2532601a6f686edbc2c6b2ec24131ff5c52f7610dd64fa4553b5464", size = 278412, upload-time = "2025-09-25T19:49:57.356Z" }, + { url = "https://files.pythonhosted.org/packages/45/2b/77424511adb11e6a99e3a00dcc7745034bee89036ad7d7e255a7e47be7d8/bcrypt-5.0.0-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:5b1589f4839a0899c146e8892efe320c0fa096568abd9b95593efac50a87cb75", size = 275486, upload-time = "2025-09-25T19:49:59.116Z" }, + { url = "https://files.pythonhosted.org/packages/43/0a/405c753f6158e0f3f14b00b462d8bca31296f7ecfc8fc8bc7919c0c7d73a/bcrypt-5.0.0-cp38-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:89042e61b5e808b67daf24a434d89bab164d4de1746b37a8d173b6b14f3db9ff", size = 277940, upload-time = "2025-09-25T19:50:00.869Z" }, + { url = "https://files.pythonhosted.org/packages/62/83/b3efc285d4aadc1fa83db385ec64dcfa1707e890eb42f03b127d66ac1b7b/bcrypt-5.0.0-cp38-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:e3cf5b2560c7b5a142286f69bde914494b6d8f901aaa71e453078388a50881c4", size = 310776, upload-time = "2025-09-25T19:50:02.393Z" }, + { url = "https://files.pythonhosted.org/packages/95/7d/47ee337dacecde6d234890fe929936cb03ebc4c3a7460854bbd9c97780b8/bcrypt-5.0.0-cp38-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:f632fd56fc4e61564f78b46a2269153122db34988e78b6be8b32d28507b7eaeb", size = 312922, upload-time = "2025-09-25T19:50:04.232Z" }, + { url = "https://files.pythonhosted.org/packages/d6/3a/43d494dfb728f55f4e1cf8fd435d50c16a2d75493225b54c8d06122523c6/bcrypt-5.0.0-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:801cad5ccb6b87d1b430f183269b94c24f248dddbbc5c1f78b6ed231743e001c", size = 341367, upload-time = "2025-09-25T19:50:05.559Z" }, + { url = "https://files.pythonhosted.org/packages/55/ab/a0727a4547e383e2e22a630e0f908113db37904f58719dc48d4622139b5c/bcrypt-5.0.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:3cf67a804fc66fc217e6914a5635000259fbbbb12e78a99488e4d5ba445a71eb", size = 359187, upload-time = "2025-09-25T19:50:06.916Z" }, + { url = "https://files.pythonhosted.org/packages/1b/bb/461f352fdca663524b4643d8b09e8435b4990f17fbf4fea6bc2a90aa0cc7/bcrypt-5.0.0-cp38-abi3-win32.whl", hash = "sha256:3abeb543874b2c0524ff40c57a4e14e5d3a66ff33fb423529c88f180fd756538", size = 153752, upload-time = "2025-09-25T19:50:08.515Z" }, + { url = "https://files.pythonhosted.org/packages/41/aa/4190e60921927b7056820291f56fc57d00d04757c8b316b2d3c0d1d6da2c/bcrypt-5.0.0-cp38-abi3-win_amd64.whl", hash = "sha256:35a77ec55b541e5e583eb3436ffbbf53b0ffa1fa16ca6782279daf95d146dcd9", size = 150881, upload-time = "2025-09-25T19:50:09.742Z" }, + { url = "https://files.pythonhosted.org/packages/54/12/cd77221719d0b39ac0b55dbd39358db1cd1246e0282e104366ebbfb8266a/bcrypt-5.0.0-cp38-abi3-win_arm64.whl", hash = "sha256:cde08734f12c6a4e28dc6755cd11d3bdfea608d93d958fffbe95a7026ebe4980", size = 144931, upload-time = "2025-09-25T19:50:11.016Z" }, + { url = "https://files.pythonhosted.org/packages/5d/ba/2af136406e1c3839aea9ecadc2f6be2bcd1eff255bd451dd39bcf302c47a/bcrypt-5.0.0-cp39-abi3-macosx_10_12_universal2.whl", hash = "sha256:0c418ca99fd47e9c59a301744d63328f17798b5947b0f791e9af3c1c499c2d0a", size = 495313, upload-time = "2025-09-25T19:50:12.309Z" }, + { url = "https://files.pythonhosted.org/packages/ac/ee/2f4985dbad090ace5ad1f7dd8ff94477fe089b5fab2040bd784a3d5f187b/bcrypt-5.0.0-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ddb4e1500f6efdd402218ffe34d040a1196c072e07929b9820f363a1fd1f4191", size = 275290, upload-time = "2025-09-25T19:50:13.673Z" }, + { url = "https://files.pythonhosted.org/packages/e4/6e/b77ade812672d15cf50842e167eead80ac3514f3beacac8902915417f8b7/bcrypt-5.0.0-cp39-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7aeef54b60ceddb6f30ee3db090351ecf0d40ec6e2abf41430997407a46d2254", size = 278253, upload-time = "2025-09-25T19:50:15.089Z" }, + { url = "https://files.pythonhosted.org/packages/36/c4/ed00ed32f1040f7990dac7115f82273e3c03da1e1a1587a778d8cea496d8/bcrypt-5.0.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:f0ce778135f60799d89c9693b9b398819d15f1921ba15fe719acb3178215a7db", size = 276084, upload-time = "2025-09-25T19:50:16.699Z" }, + { url = "https://files.pythonhosted.org/packages/e7/c4/fa6e16145e145e87f1fa351bbd54b429354fd72145cd3d4e0c5157cf4c70/bcrypt-5.0.0-cp39-abi3-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:a71f70ee269671460b37a449f5ff26982a6f2ba493b3eabdd687b4bf35f875ac", size = 297185, upload-time = "2025-09-25T19:50:18.525Z" }, + { url = "https://files.pythonhosted.org/packages/24/b4/11f8a31d8b67cca3371e046db49baa7c0594d71eb40ac8121e2fc0888db0/bcrypt-5.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:f8429e1c410b4073944f03bd778a9e066e7fad723564a52ff91841d278dfc822", size = 278656, upload-time = "2025-09-25T19:50:19.809Z" }, + { url = "https://files.pythonhosted.org/packages/ac/31/79f11865f8078e192847d2cb526e3fa27c200933c982c5b2869720fa5fce/bcrypt-5.0.0-cp39-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:edfcdcedd0d0f05850c52ba3127b1fce70b9f89e0fe5ff16517df7e81fa3cbb8", size = 275662, upload-time = "2025-09-25T19:50:21.567Z" }, + { url = "https://files.pythonhosted.org/packages/d4/8d/5e43d9584b3b3591a6f9b68f755a4da879a59712981ef5ad2a0ac1379f7a/bcrypt-5.0.0-cp39-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:611f0a17aa4a25a69362dcc299fda5c8a3d4f160e2abb3831041feb77393a14a", size = 278240, upload-time = "2025-09-25T19:50:23.305Z" }, + { url = "https://files.pythonhosted.org/packages/89/48/44590e3fc158620f680a978aafe8f87a4c4320da81ed11552f0323aa9a57/bcrypt-5.0.0-cp39-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:db99dca3b1fdc3db87d7c57eac0c82281242d1eabf19dcb8a6b10eb29a2e72d1", size = 311152, upload-time = "2025-09-25T19:50:24.597Z" }, + { url = "https://files.pythonhosted.org/packages/5f/85/e4fbfc46f14f47b0d20493669a625da5827d07e8a88ee460af6cd9768b44/bcrypt-5.0.0-cp39-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:5feebf85a9cefda32966d8171f5db7e3ba964b77fdfe31919622256f80f9cf42", size = 313284, upload-time = "2025-09-25T19:50:26.268Z" }, + { url = "https://files.pythonhosted.org/packages/25/ae/479f81d3f4594456a01ea2f05b132a519eff9ab5768a70430fa1132384b1/bcrypt-5.0.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:3ca8a166b1140436e058298a34d88032ab62f15aae1c598580333dc21d27ef10", size = 341643, upload-time = "2025-09-25T19:50:28.02Z" }, + { url = "https://files.pythonhosted.org/packages/df/d2/36a086dee1473b14276cd6ea7f61aef3b2648710b5d7f1c9e032c29b859f/bcrypt-5.0.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:61afc381250c3182d9078551e3ac3a41da14154fbff647ddf52a769f588c4172", size = 359698, upload-time = "2025-09-25T19:50:31.347Z" }, + { url = "https://files.pythonhosted.org/packages/c0/f6/688d2cd64bfd0b14d805ddb8a565e11ca1fb0fd6817175d58b10052b6d88/bcrypt-5.0.0-cp39-abi3-win32.whl", hash = "sha256:64d7ce196203e468c457c37ec22390f1a61c85c6f0b8160fd752940ccfb3a683", size = 153725, upload-time = "2025-09-25T19:50:34.384Z" }, + { url = "https://files.pythonhosted.org/packages/9f/b9/9d9a641194a730bda138b3dfe53f584d61c58cd5230e37566e83ec2ffa0d/bcrypt-5.0.0-cp39-abi3-win_amd64.whl", hash = "sha256:64ee8434b0da054d830fa8e89e1c8bf30061d539044a39524ff7dec90481e5c2", size = 150912, upload-time = "2025-09-25T19:50:35.69Z" }, + { url = "https://files.pythonhosted.org/packages/27/44/d2ef5e87509158ad2187f4dd0852df80695bb1ee0cfe0a684727b01a69e0/bcrypt-5.0.0-cp39-abi3-win_arm64.whl", hash = "sha256:f2347d3534e76bf50bca5500989d6c1d05ed64b440408057a37673282c654927", size = 144953, upload-time = "2025-09-25T19:50:37.32Z" }, +] + +[[package]] +name = "certifi" +version = "2026.2.25" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/af/2d/7bf41579a8986e348fa033a31cdd0e4121114f6bce2457e8876010b092dd/certifi-2026.2.25.tar.gz", hash = "sha256:e887ab5cee78ea814d3472169153c2d12cd43b14bd03329a39a9c6e2e80bfba7", size = 155029, upload-time = "2026-02-25T02:54:17.342Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9a/3c/c17fb3ca2d9c3acff52e30b309f538586f9f5b9c9cf454f3845fc9af4881/certifi-2026.2.25-py3-none-any.whl", hash = "sha256:027692e4402ad994f1c42e52a4997a9763c646b73e4096e4d5d6db8af1d6f0fa", size = 153684, upload-time = "2026-02-25T02:54:15.766Z" }, +] + +[[package]] +name = "cffi" +version = "2.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pycparser", marker = "implementation_name != 'PyPy'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8453301356628e8147c79dbb825bcbc73dc7401f9846/cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529", size = 523588, upload-time = "2025-09-08T23:24:04.541Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ea/47/4f61023ea636104d4f16ab488e268b93008c3d0bb76893b1b31db1f96802/cffi-2.0.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:6d02d6655b0e54f54c4ef0b94eb6be0607b70853c45ce98bd278dc7de718be5d", size = 185271, upload-time = "2025-09-08T23:22:44.795Z" }, + { url = "https://files.pythonhosted.org/packages/df/a2/781b623f57358e360d62cdd7a8c681f074a71d445418a776eef0aadb4ab4/cffi-2.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8eca2a813c1cb7ad4fb74d368c2ffbbb4789d377ee5bb8df98373c2cc0dee76c", size = 181048, upload-time = "2025-09-08T23:22:45.938Z" }, + { url = "https://files.pythonhosted.org/packages/ff/df/a4f0fbd47331ceeba3d37c2e51e9dfc9722498becbeec2bd8bc856c9538a/cffi-2.0.0-cp312-cp312-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:21d1152871b019407d8ac3985f6775c079416c282e431a4da6afe7aefd2bccbe", size = 212529, upload-time = "2025-09-08T23:22:47.349Z" }, + { url = "https://files.pythonhosted.org/packages/d5/72/12b5f8d3865bf0f87cf1404d8c374e7487dcf097a1c91c436e72e6badd83/cffi-2.0.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b21e08af67b8a103c71a250401c78d5e0893beff75e28c53c98f4de42f774062", size = 220097, upload-time = "2025-09-08T23:22:48.677Z" }, + { url = "https://files.pythonhosted.org/packages/c2/95/7a135d52a50dfa7c882ab0ac17e8dc11cec9d55d2c18dda414c051c5e69e/cffi-2.0.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:1e3a615586f05fc4065a8b22b8152f0c1b00cdbc60596d187c2a74f9e3036e4e", size = 207983, upload-time = "2025-09-08T23:22:50.06Z" }, + { url = "https://files.pythonhosted.org/packages/3a/c8/15cb9ada8895957ea171c62dc78ff3e99159ee7adb13c0123c001a2546c1/cffi-2.0.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:81afed14892743bbe14dacb9e36d9e0e504cd204e0b165062c488942b9718037", size = 206519, upload-time = "2025-09-08T23:22:51.364Z" }, + { url = "https://files.pythonhosted.org/packages/78/2d/7fa73dfa841b5ac06c7b8855cfc18622132e365f5b81d02230333ff26e9e/cffi-2.0.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3e17ed538242334bf70832644a32a7aae3d83b57567f9fd60a26257e992b79ba", size = 219572, upload-time = "2025-09-08T23:22:52.902Z" }, + { url = "https://files.pythonhosted.org/packages/07/e0/267e57e387b4ca276b90f0434ff88b2c2241ad72b16d31836adddfd6031b/cffi-2.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3925dd22fa2b7699ed2617149842d2e6adde22b262fcbfada50e3d195e4b3a94", size = 222963, upload-time = "2025-09-08T23:22:54.518Z" }, + { url = "https://files.pythonhosted.org/packages/b6/75/1f2747525e06f53efbd878f4d03bac5b859cbc11c633d0fb81432d98a795/cffi-2.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:2c8f814d84194c9ea681642fd164267891702542f028a15fc97d4674b6206187", size = 221361, upload-time = "2025-09-08T23:22:55.867Z" }, + { url = "https://files.pythonhosted.org/packages/7b/2b/2b6435f76bfeb6bbf055596976da087377ede68df465419d192acf00c437/cffi-2.0.0-cp312-cp312-win32.whl", hash = "sha256:da902562c3e9c550df360bfa53c035b2f241fed6d9aef119048073680ace4a18", size = 172932, upload-time = "2025-09-08T23:22:57.188Z" }, + { url = "https://files.pythonhosted.org/packages/f8/ed/13bd4418627013bec4ed6e54283b1959cf6db888048c7cf4b4c3b5b36002/cffi-2.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:da68248800ad6320861f129cd9c1bf96ca849a2771a59e0344e88681905916f5", size = 183557, upload-time = "2025-09-08T23:22:58.351Z" }, + { url = "https://files.pythonhosted.org/packages/95/31/9f7f93ad2f8eff1dbc1c3656d7ca5bfd8fb52c9d786b4dcf19b2d02217fa/cffi-2.0.0-cp312-cp312-win_arm64.whl", hash = "sha256:4671d9dd5ec934cb9a73e7ee9676f9362aba54f7f34910956b84d727b0d73fb6", size = 177762, upload-time = "2025-09-08T23:22:59.668Z" }, + { url = "https://files.pythonhosted.org/packages/4b/8d/a0a47a0c9e413a658623d014e91e74a50cdd2c423f7ccfd44086ef767f90/cffi-2.0.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:00bdf7acc5f795150faa6957054fbbca2439db2f775ce831222b66f192f03beb", size = 185230, upload-time = "2025-09-08T23:23:00.879Z" }, + { url = "https://files.pythonhosted.org/packages/4a/d2/a6c0296814556c68ee32009d9c2ad4f85f2707cdecfd7727951ec228005d/cffi-2.0.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:45d5e886156860dc35862657e1494b9bae8dfa63bf56796f2fb56e1679fc0bca", size = 181043, upload-time = "2025-09-08T23:23:02.231Z" }, + { url = "https://files.pythonhosted.org/packages/b0/1e/d22cc63332bd59b06481ceaac49d6c507598642e2230f201649058a7e704/cffi-2.0.0-cp313-cp313-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:07b271772c100085dd28b74fa0cd81c8fb1a3ba18b21e03d7c27f3436a10606b", size = 212446, upload-time = "2025-09-08T23:23:03.472Z" }, + { url = "https://files.pythonhosted.org/packages/a9/f5/a2c23eb03b61a0b8747f211eb716446c826ad66818ddc7810cc2cc19b3f2/cffi-2.0.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d48a880098c96020b02d5a1f7d9251308510ce8858940e6fa99ece33f610838b", size = 220101, upload-time = "2025-09-08T23:23:04.792Z" }, + { url = "https://files.pythonhosted.org/packages/f2/7f/e6647792fc5850d634695bc0e6ab4111ae88e89981d35ac269956605feba/cffi-2.0.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:f93fd8e5c8c0a4aa1f424d6173f14a892044054871c771f8566e4008eaa359d2", size = 207948, upload-time = "2025-09-08T23:23:06.127Z" }, + { url = "https://files.pythonhosted.org/packages/cb/1e/a5a1bd6f1fb30f22573f76533de12a00bf274abcdc55c8edab639078abb6/cffi-2.0.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:dd4f05f54a52fb558f1ba9f528228066954fee3ebe629fc1660d874d040ae5a3", size = 206422, upload-time = "2025-09-08T23:23:07.753Z" }, + { url = "https://files.pythonhosted.org/packages/98/df/0a1755e750013a2081e863e7cd37e0cdd02664372c754e5560099eb7aa44/cffi-2.0.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c8d3b5532fc71b7a77c09192b4a5a200ea992702734a2e9279a37f2478236f26", size = 219499, upload-time = "2025-09-08T23:23:09.648Z" }, + { url = "https://files.pythonhosted.org/packages/50/e1/a969e687fcf9ea58e6e2a928ad5e2dd88cc12f6f0ab477e9971f2309b57c/cffi-2.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d9b29c1f0ae438d5ee9acb31cadee00a58c46cc9c0b2f9038c6b0b3470877a8c", size = 222928, upload-time = "2025-09-08T23:23:10.928Z" }, + { url = "https://files.pythonhosted.org/packages/36/54/0362578dd2c9e557a28ac77698ed67323ed5b9775ca9d3fe73fe191bb5d8/cffi-2.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6d50360be4546678fc1b79ffe7a66265e28667840010348dd69a314145807a1b", size = 221302, upload-time = "2025-09-08T23:23:12.42Z" }, + { url = "https://files.pythonhosted.org/packages/eb/6d/bf9bda840d5f1dfdbf0feca87fbdb64a918a69bca42cfa0ba7b137c48cb8/cffi-2.0.0-cp313-cp313-win32.whl", hash = "sha256:74a03b9698e198d47562765773b4a8309919089150a0bb17d829ad7b44b60d27", size = 172909, upload-time = "2025-09-08T23:23:14.32Z" }, + { url = "https://files.pythonhosted.org/packages/37/18/6519e1ee6f5a1e579e04b9ddb6f1676c17368a7aba48299c3759bbc3c8b3/cffi-2.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:19f705ada2530c1167abacb171925dd886168931e0a7b78f5bffcae5c6b5be75", size = 183402, upload-time = "2025-09-08T23:23:15.535Z" }, + { url = "https://files.pythonhosted.org/packages/cb/0e/02ceeec9a7d6ee63bb596121c2c8e9b3a9e150936f4fbef6ca1943e6137c/cffi-2.0.0-cp313-cp313-win_arm64.whl", hash = "sha256:256f80b80ca3853f90c21b23ee78cd008713787b1b1e93eae9f3d6a7134abd91", size = 177780, upload-time = "2025-09-08T23:23:16.761Z" }, +] + +[[package]] +name = "charset-normalizer" +version = "3.4.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e7/a1/67fe25fac3c7642725500a3f6cfe5821ad557c3abb11c9d20d12c7008d3e/charset_normalizer-3.4.7.tar.gz", hash = "sha256:ae89db9e5f98a11a4bf50407d4363e7b09b31e55bc117b4f7d80aab97ba009e5", size = 144271, upload-time = "2026-04-02T09:28:39.342Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0c/eb/4fc8d0a7110eb5fc9cc161723a34a8a6c200ce3b4fbf681bc86feee22308/charset_normalizer-3.4.7-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:eca9705049ad3c7345d574e3510665cb2cf844c2f2dcfe675332677f081cbd46", size = 311328, upload-time = "2026-04-02T09:26:24.331Z" }, + { url = "https://files.pythonhosted.org/packages/f8/e3/0fadc706008ac9d7b9b5be6dc767c05f9d3e5df51744ce4cc9605de7b9f4/charset_normalizer-3.4.7-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6178f72c5508bfc5fd446a5905e698c6212932f25bcdd4b47a757a50605a90e2", size = 208061, upload-time = "2026-04-02T09:26:25.568Z" }, + { url = "https://files.pythonhosted.org/packages/42/f0/3dd1045c47f4a4604df85ec18ad093912ae1344ac706993aff91d38773a2/charset_normalizer-3.4.7-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e1421b502d83040e6d7fb2fb18dff63957f720da3d77b2fbd3187ceb63755d7b", size = 229031, upload-time = "2026-04-02T09:26:26.865Z" }, + { url = "https://files.pythonhosted.org/packages/dc/67/675a46eb016118a2fbde5a277a5d15f4f69d5f3f5f338e5ee2f8948fcf43/charset_normalizer-3.4.7-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:edac0f1ab77644605be2cbba52e6b7f630731fc42b34cb0f634be1a6eface56a", size = 225239, upload-time = "2026-04-02T09:26:28.044Z" }, + { url = "https://files.pythonhosted.org/packages/4b/f8/d0118a2f5f23b02cd166fa385c60f9b0d4f9194f574e2b31cef350ad7223/charset_normalizer-3.4.7-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5649fd1c7bade02f320a462fdefd0b4bd3ce036065836d4f42e0de958038e116", size = 216589, upload-time = "2026-04-02T09:26:29.239Z" }, + { url = "https://files.pythonhosted.org/packages/b1/f1/6d2b0b261b6c4ceef0fcb0d17a01cc5bc53586c2d4796fa04b5c540bc13d/charset_normalizer-3.4.7-cp312-cp312-manylinux_2_31_armv7l.whl", hash = "sha256:203104ed3e428044fd943bc4bf45fa73c0730391f9621e37fe39ecf477b128cb", size = 202733, upload-time = "2026-04-02T09:26:30.5Z" }, + { url = "https://files.pythonhosted.org/packages/6f/c0/7b1f943f7e87cc3db9626ba17807d042c38645f0a1d4415c7a14afb5591f/charset_normalizer-3.4.7-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:298930cec56029e05497a76988377cbd7457ba864beeea92ad7e844fe74cd1f1", size = 212652, upload-time = "2026-04-02T09:26:31.709Z" }, + { url = "https://files.pythonhosted.org/packages/38/dd/5a9ab159fe45c6e72079398f277b7d2b523e7f716acc489726115a910097/charset_normalizer-3.4.7-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:708838739abf24b2ceb208d0e22403dd018faeef86ddac04319a62ae884c4f15", size = 211229, upload-time = "2026-04-02T09:26:33.282Z" }, + { url = "https://files.pythonhosted.org/packages/d5/ff/531a1cad5ca855d1c1a8b69cb71abfd6d85c0291580146fda7c82857caa1/charset_normalizer-3.4.7-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:0f7eb884681e3938906ed0434f20c63046eacd0111c4ba96f27b76084cd679f5", size = 203552, upload-time = "2026-04-02T09:26:34.845Z" }, + { url = "https://files.pythonhosted.org/packages/c1/4c/a5fb52d528a8ca41f7598cb619409ece30a169fbdf9cdce592e53b46c3a6/charset_normalizer-3.4.7-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:4dc1e73c36828f982bfe79fadf5919923f8a6f4df2860804db9a98c48824ce8d", size = 230806, upload-time = "2026-04-02T09:26:36.152Z" }, + { url = "https://files.pythonhosted.org/packages/59/7a/071feed8124111a32b316b33ae4de83d36923039ef8cf48120266844285b/charset_normalizer-3.4.7-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:aed52fea0513bac0ccde438c188c8a471c4e0f457c2dd20cdbf6ea7a450046c7", size = 212316, upload-time = "2026-04-02T09:26:37.672Z" }, + { url = "https://files.pythonhosted.org/packages/fd/35/f7dba3994312d7ba508e041eaac39a36b120f32d4c8662b8814dab876431/charset_normalizer-3.4.7-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:fea24543955a6a729c45a73fe90e08c743f0b3334bbf3201e6c4bc1b0c7fa464", size = 227274, upload-time = "2026-04-02T09:26:38.93Z" }, + { url = "https://files.pythonhosted.org/packages/8a/2d/a572df5c9204ab7688ec1edc895a73ebded3b023bb07364710b05dd1c9be/charset_normalizer-3.4.7-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:bb6d88045545b26da47aa879dd4a89a71d1dce0f0e549b1abcb31dfe4a8eac49", size = 218468, upload-time = "2026-04-02T09:26:40.17Z" }, + { url = "https://files.pythonhosted.org/packages/86/eb/890922a8b03a568ca2f336c36585a4713c55d4d67bf0f0c78924be6315ca/charset_normalizer-3.4.7-cp312-cp312-win32.whl", hash = "sha256:2257141f39fe65a3fdf38aeccae4b953e5f3b3324f4ff0daf9f15b8518666a2c", size = 148460, upload-time = "2026-04-02T09:26:41.416Z" }, + { url = "https://files.pythonhosted.org/packages/35/d9/0e7dffa06c5ab081f75b1b786f0aefc88365825dfcd0ac544bdb7b2b6853/charset_normalizer-3.4.7-cp312-cp312-win_amd64.whl", hash = "sha256:5ed6ab538499c8644b8a3e18debabcd7ce684f3fa91cf867521a7a0279cab2d6", size = 159330, upload-time = "2026-04-02T09:26:42.554Z" }, + { url = "https://files.pythonhosted.org/packages/9e/5d/481bcc2a7c88ea6b0878c299547843b2521ccbc40980cb406267088bc701/charset_normalizer-3.4.7-cp312-cp312-win_arm64.whl", hash = "sha256:56be790f86bfb2c98fb742ce566dfb4816e5a83384616ab59c49e0604d49c51d", size = 147828, upload-time = "2026-04-02T09:26:44.075Z" }, + { url = "https://files.pythonhosted.org/packages/c1/3b/66777e39d3ae1ddc77ee606be4ec6d8cbd4c801f65e5a1b6f2b11b8346dd/charset_normalizer-3.4.7-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:f496c9c3cc02230093d8330875c4c3cdfc3b73612a5fd921c65d39cbcef08063", size = 309627, upload-time = "2026-04-02T09:26:45.198Z" }, + { url = "https://files.pythonhosted.org/packages/2e/4e/b7f84e617b4854ade48a1b7915c8ccfadeba444d2a18c291f696e37f0d3b/charset_normalizer-3.4.7-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0ea948db76d31190bf08bd371623927ee1339d5f2a0b4b1b4a4439a65298703c", size = 207008, upload-time = "2026-04-02T09:26:46.824Z" }, + { url = "https://files.pythonhosted.org/packages/c4/bb/ec73c0257c9e11b268f018f068f5d00aa0ef8c8b09f7753ebd5f2880e248/charset_normalizer-3.4.7-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a277ab8928b9f299723bc1a2dabb1265911b1a76341f90a510368ca44ad9ab66", size = 228303, upload-time = "2026-04-02T09:26:48.397Z" }, + { url = "https://files.pythonhosted.org/packages/85/fb/32d1f5033484494619f701e719429c69b766bfc4dbc61aa9e9c8c166528b/charset_normalizer-3.4.7-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3bec022aec2c514d9cf199522a802bd007cd588ab17ab2525f20f9c34d067c18", size = 224282, upload-time = "2026-04-02T09:26:49.684Z" }, + { url = "https://files.pythonhosted.org/packages/fa/07/330e3a0dda4c404d6da83b327270906e9654a24f6c546dc886a0eb0ffb23/charset_normalizer-3.4.7-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e044c39e41b92c845bc815e5ae4230804e8e7bc29e399b0437d64222d92809dd", size = 215595, upload-time = "2026-04-02T09:26:50.915Z" }, + { url = "https://files.pythonhosted.org/packages/e3/7c/fc890655786e423f02556e0216d4b8c6bcb6bdfa890160dc66bf52dee468/charset_normalizer-3.4.7-cp313-cp313-manylinux_2_31_armv7l.whl", hash = "sha256:f495a1652cf3fbab2eb0639776dad966c2fb874d79d87ca07f9d5f059b8bd215", size = 201986, upload-time = "2026-04-02T09:26:52.197Z" }, + { url = "https://files.pythonhosted.org/packages/d8/97/bfb18b3db2aed3b90cf54dc292ad79fdd5ad65c4eae454099475cbeadd0d/charset_normalizer-3.4.7-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e712b419df8ba5e42b226c510472b37bd57b38e897d3eca5e8cfd410a29fa859", size = 211711, upload-time = "2026-04-02T09:26:53.49Z" }, + { url = "https://files.pythonhosted.org/packages/6f/a5/a581c13798546a7fd557c82614a5c65a13df2157e9ad6373166d2a3e645d/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:7804338df6fcc08105c7745f1502ba68d900f45fd770d5bdd5288ddccb8a42d8", size = 210036, upload-time = "2026-04-02T09:26:54.975Z" }, + { url = "https://files.pythonhosted.org/packages/8c/bf/b3ab5bcb478e4193d517644b0fb2bf5497fbceeaa7a1bc0f4d5b50953861/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:481551899c856c704d58119b5025793fa6730adda3571971af568f66d2424bb5", size = 202998, upload-time = "2026-04-02T09:26:56.303Z" }, + { url = "https://files.pythonhosted.org/packages/e7/4e/23efd79b65d314fa320ec6017b4b5834d5c12a58ba4610aa353af2e2f577/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:f59099f9b66f0d7145115e6f80dd8b1d847176df89b234a5a6b3f00437aa0832", size = 230056, upload-time = "2026-04-02T09:26:57.554Z" }, + { url = "https://files.pythonhosted.org/packages/b9/9f/1e1941bc3f0e01df116e68dc37a55c4d249df5e6fa77f008841aef68264f/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:f59ad4c0e8f6bba240a9bb85504faa1ab438237199d4cce5f622761507b8f6a6", size = 211537, upload-time = "2026-04-02T09:26:58.843Z" }, + { url = "https://files.pythonhosted.org/packages/80/0f/088cbb3020d44428964a6c97fe1edfb1b9550396bf6d278330281e8b709c/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:3dedcc22d73ec993f42055eff4fcfed9318d1eeb9a6606c55892a26964964e48", size = 226176, upload-time = "2026-04-02T09:27:00.437Z" }, + { url = "https://files.pythonhosted.org/packages/6a/9f/130394f9bbe06f4f63e22641d32fc9b202b7e251c9aef4db044324dac493/charset_normalizer-3.4.7-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:64f02c6841d7d83f832cd97ccf8eb8a906d06eb95d5276069175c696b024b60a", size = 217723, upload-time = "2026-04-02T09:27:02.021Z" }, + { url = "https://files.pythonhosted.org/packages/73/55/c469897448a06e49f8fa03f6caae97074fde823f432a98f979cc42b90e69/charset_normalizer-3.4.7-cp313-cp313-win32.whl", hash = "sha256:4042d5c8f957e15221d423ba781e85d553722fc4113f523f2feb7b188cc34c5e", size = 148085, upload-time = "2026-04-02T09:27:03.192Z" }, + { url = "https://files.pythonhosted.org/packages/5d/78/1b74c5bbb3f99b77a1715c91b3e0b5bdb6fe302d95ace4f5b1bec37b0167/charset_normalizer-3.4.7-cp313-cp313-win_amd64.whl", hash = "sha256:3946fa46a0cf3e4c8cb1cc52f56bb536310d34f25f01ca9b6c16afa767dab110", size = 158819, upload-time = "2026-04-02T09:27:04.454Z" }, + { url = "https://files.pythonhosted.org/packages/68/86/46bd42279d323deb8687c4a5a811fd548cb7d1de10cf6535d099877a9a9f/charset_normalizer-3.4.7-cp313-cp313-win_arm64.whl", hash = "sha256:80d04837f55fc81da168b98de4f4b797ef007fc8a79ab71c6ec9bc4dd662b15b", size = 147915, upload-time = "2026-04-02T09:27:05.971Z" }, + { url = "https://files.pythonhosted.org/packages/db/8f/61959034484a4a7c527811f4721e75d02d653a35afb0b6054474d8185d4c/charset_normalizer-3.4.7-py3-none-any.whl", hash = "sha256:3dce51d0f5e7951f8bb4900c257dad282f49190fdbebecd4ba99bcc41fef404d", size = 61958, upload-time = "2026-04-02T09:28:37.794Z" }, +] + +[[package]] +name = "click" +version = "8.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/57/75/31212c6bf2503fdf920d87fee5d7a86a2e3bcf444984126f13d8e4016804/click-8.3.2.tar.gz", hash = "sha256:14162b8b3b3550a7d479eafa77dfd3c38d9dc8951f6f69c78913a8f9a7540fd5", size = 302856, upload-time = "2026-04-03T19:14:45.118Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e4/20/71885d8b97d4f3dde17b1fdb92dbd4908b00541c5a3379787137285f602e/click-8.3.2-py3-none-any.whl", hash = "sha256:1924d2c27c5653561cd2cae4548d1406039cb79b858b747cfea24924bbc1616d", size = 108379, upload-time = "2026-04-03T19:14:43.505Z" }, +] + +[[package]] +name = "cloudpickle" +version = "3.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, +] + +[[package]] +name = "colorama" +version = "0.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, +] + +[[package]] +name = "colorlog" +version = "6.10.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a2/61/f083b5ac52e505dfc1c624eafbf8c7589a0d7f32daa398d2e7590efa5fda/colorlog-6.10.1.tar.gz", hash = "sha256:eb4ae5cb65fe7fec7773c2306061a8e63e02efc2c72eba9d27b0fa23c94f1321", size = 17162, upload-time = "2025-10-16T16:14:11.978Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6d/c1/e419ef3723a074172b68aaa89c9f3de486ed4c2399e2dbd8113a4fdcaf9e/colorlog-6.10.1-py3-none-any.whl", hash = "sha256:2d7e8348291948af66122cff006c9f8da6255d224e7cf8e37d8de2df3bad8c9c", size = 11743, upload-time = "2025-10-16T16:14:10.512Z" }, +] + +[[package]] +name = "cryptography" +version = "46.0.7" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cffi", marker = "platform_python_implementation != 'PyPy'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/47/93/ac8f3d5ff04d54bc814e961a43ae5b0b146154c89c61b47bb07557679b18/cryptography-46.0.7.tar.gz", hash = "sha256:e4cfd68c5f3e0bfdad0d38e023239b96a2fe84146481852dffbcca442c245aa5", size = 750652, upload-time = "2026-04-08T01:57:54.692Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0b/5d/4a8f770695d73be252331e60e526291e3df0c9b27556a90a6b47bccca4c2/cryptography-46.0.7-cp311-abi3-macosx_10_9_universal2.whl", hash = "sha256:ea42cbe97209df307fdc3b155f1b6fa2577c0defa8f1f7d3be7d31d189108ad4", size = 7179869, upload-time = "2026-04-08T01:56:17.157Z" }, + { url = "https://files.pythonhosted.org/packages/5f/45/6d80dc379b0bbc1f9d1e429f42e4cb9e1d319c7a8201beffd967c516ea01/cryptography-46.0.7-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b36a4695e29fe69215d75960b22577197aca3f7a25b9cf9d165dcfe9d80bc325", size = 4275492, upload-time = "2026-04-08T01:56:19.36Z" }, + { url = "https://files.pythonhosted.org/packages/4a/9a/1765afe9f572e239c3469f2cb429f3ba7b31878c893b246b4b2994ffe2fe/cryptography-46.0.7-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5ad9ef796328c5e3c4ceed237a183f5d41d21150f972455a9d926593a1dcb308", size = 4426670, upload-time = "2026-04-08T01:56:21.415Z" }, + { url = "https://files.pythonhosted.org/packages/8f/3e/af9246aaf23cd4ee060699adab1e47ced3f5f7e7a8ffdd339f817b446462/cryptography-46.0.7-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:73510b83623e080a2c35c62c15298096e2a5dc8d51c3b4e1740211839d0dea77", size = 4280275, upload-time = "2026-04-08T01:56:23.539Z" }, + { url = "https://files.pythonhosted.org/packages/0f/54/6bbbfc5efe86f9d71041827b793c24811a017c6ac0fd12883e4caa86b8ed/cryptography-46.0.7-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:cbd5fb06b62bd0721e1170273d3f4d5a277044c47ca27ee257025146c34cbdd1", size = 4928402, upload-time = "2026-04-08T01:56:25.624Z" }, + { url = "https://files.pythonhosted.org/packages/2d/cf/054b9d8220f81509939599c8bdbc0c408dbd2bdd41688616a20731371fe0/cryptography-46.0.7-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:420b1e4109cc95f0e5700eed79908cef9268265c773d3a66f7af1eef53d409ef", size = 4459985, upload-time = "2026-04-08T01:56:27.309Z" }, + { url = "https://files.pythonhosted.org/packages/f9/46/4e4e9c6040fb01c7467d47217d2f882daddeb8828f7df800cb806d8a2288/cryptography-46.0.7-cp311-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:24402210aa54baae71d99441d15bb5a1919c195398a87b563df84468160a65de", size = 3990652, upload-time = "2026-04-08T01:56:29.095Z" }, + { url = "https://files.pythonhosted.org/packages/36/5f/313586c3be5a2fbe87e4c9a254207b860155a8e1f3cca99f9910008e7d08/cryptography-46.0.7-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:8a469028a86f12eb7d2fe97162d0634026d92a21f3ae0ac87ed1c4a447886c83", size = 4279805, upload-time = "2026-04-08T01:56:30.928Z" }, + { url = "https://files.pythonhosted.org/packages/69/33/60dfc4595f334a2082749673386a4d05e4f0cf4df8248e63b2c3437585f2/cryptography-46.0.7-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:9694078c5d44c157ef3162e3bf3946510b857df5a3955458381d1c7cfc143ddb", size = 4892883, upload-time = "2026-04-08T01:56:32.614Z" }, + { url = "https://files.pythonhosted.org/packages/c7/0b/333ddab4270c4f5b972f980adef4faa66951a4aaf646ca067af597f15563/cryptography-46.0.7-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:42a1e5f98abb6391717978baf9f90dc28a743b7d9be7f0751a6f56a75d14065b", size = 4459756, upload-time = "2026-04-08T01:56:34.306Z" }, + { url = "https://files.pythonhosted.org/packages/d2/14/633913398b43b75f1234834170947957c6b623d1701ffc7a9600da907e89/cryptography-46.0.7-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:91bbcb08347344f810cbe49065914fe048949648f6bd5c2519f34619142bbe85", size = 4410244, upload-time = "2026-04-08T01:56:35.977Z" }, + { url = "https://files.pythonhosted.org/packages/10/f2/19ceb3b3dc14009373432af0c13f46aa08e3ce334ec6eff13492e1812ccd/cryptography-46.0.7-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:5d1c02a14ceb9148cc7816249f64f623fbfee39e8c03b3650d842ad3f34d637e", size = 4674868, upload-time = "2026-04-08T01:56:38.034Z" }, + { url = "https://files.pythonhosted.org/packages/1a/bb/a5c213c19ee94b15dfccc48f363738633a493812687f5567addbcbba9f6f/cryptography-46.0.7-cp311-abi3-win32.whl", hash = "sha256:d23c8ca48e44ee015cd0a54aeccdf9f09004eba9fc96f38c911011d9ff1bd457", size = 3026504, upload-time = "2026-04-08T01:56:39.666Z" }, + { url = "https://files.pythonhosted.org/packages/2b/02/7788f9fefa1d060ca68717c3901ae7fffa21ee087a90b7f23c7a603c32ae/cryptography-46.0.7-cp311-abi3-win_amd64.whl", hash = "sha256:397655da831414d165029da9bc483bed2fe0e75dde6a1523ec2fe63f3c46046b", size = 3488363, upload-time = "2026-04-08T01:56:41.893Z" }, + { url = "https://files.pythonhosted.org/packages/a7/7f/cd42fc3614386bc0c12f0cb3c4ae1fc2bbca5c9662dfed031514911d513d/cryptography-46.0.7-cp38-abi3-macosx_10_9_universal2.whl", hash = "sha256:462ad5cb1c148a22b2e3bcc5ad52504dff325d17daf5df8d88c17dda1f75f2a4", size = 7165618, upload-time = "2026-04-08T01:57:10.645Z" }, + { url = "https://files.pythonhosted.org/packages/a5/d0/36a49f0262d2319139d2829f773f1b97ef8aef7f97e6e5bd21455e5a8fb5/cryptography-46.0.7-cp38-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:84d4cced91f0f159a7ddacad249cc077e63195c36aac40b4150e7a57e84fffe7", size = 4270628, upload-time = "2026-04-08T01:57:12.885Z" }, + { url = "https://files.pythonhosted.org/packages/8a/6c/1a42450f464dda6ffbe578a911f773e54dd48c10f9895a23a7e88b3e7db5/cryptography-46.0.7-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:128c5edfe5e5938b86b03941e94fac9ee793a94452ad1365c9fc3f4f62216832", size = 4415405, upload-time = "2026-04-08T01:57:14.923Z" }, + { url = "https://files.pythonhosted.org/packages/9a/92/4ed714dbe93a066dc1f4b4581a464d2d7dbec9046f7c8b7016f5286329e2/cryptography-46.0.7-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:5e51be372b26ef4ba3de3c167cd3d1022934bc838ae9eaad7e644986d2a3d163", size = 4272715, upload-time = "2026-04-08T01:57:16.638Z" }, + { url = "https://files.pythonhosted.org/packages/b7/e6/a26b84096eddd51494bba19111f8fffe976f6a09f132706f8f1bf03f51f7/cryptography-46.0.7-cp38-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:cdf1a610ef82abb396451862739e3fc93b071c844399e15b90726ef7470eeaf2", size = 4918400, upload-time = "2026-04-08T01:57:19.021Z" }, + { url = "https://files.pythonhosted.org/packages/c7/08/ffd537b605568a148543ac3c2b239708ae0bd635064bab41359252ef88ed/cryptography-46.0.7-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:1d25aee46d0c6f1a501adcddb2d2fee4b979381346a78558ed13e50aa8a59067", size = 4450634, upload-time = "2026-04-08T01:57:21.185Z" }, + { url = "https://files.pythonhosted.org/packages/16/01/0cd51dd86ab5b9befe0d031e276510491976c3a80e9f6e31810cce46c4ad/cryptography-46.0.7-cp38-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:cdfbe22376065ffcf8be74dc9a909f032df19bc58a699456a21712d6e5eabfd0", size = 3985233, upload-time = "2026-04-08T01:57:22.862Z" }, + { url = "https://files.pythonhosted.org/packages/92/49/819d6ed3a7d9349c2939f81b500a738cb733ab62fbecdbc1e38e83d45e12/cryptography-46.0.7-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:abad9dac36cbf55de6eb49badd4016806b3165d396f64925bf2999bcb67837ba", size = 4271955, upload-time = "2026-04-08T01:57:24.814Z" }, + { url = "https://files.pythonhosted.org/packages/80/07/ad9b3c56ebb95ed2473d46df0847357e01583f4c52a85754d1a55e29e4d0/cryptography-46.0.7-cp38-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:935ce7e3cfdb53e3536119a542b839bb94ec1ad081013e9ab9b7cfd478b05006", size = 4879888, upload-time = "2026-04-08T01:57:26.88Z" }, + { url = "https://files.pythonhosted.org/packages/b8/c7/201d3d58f30c4c2bdbe9b03844c291feb77c20511cc3586daf7edc12a47b/cryptography-46.0.7-cp38-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:35719dc79d4730d30f1c2b6474bd6acda36ae2dfae1e3c16f2051f215df33ce0", size = 4449961, upload-time = "2026-04-08T01:57:29.068Z" }, + { url = "https://files.pythonhosted.org/packages/a5/ef/649750cbf96f3033c3c976e112265c33906f8e462291a33d77f90356548c/cryptography-46.0.7-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:7bbc6ccf49d05ac8f7d7b5e2e2c33830d4fe2061def88210a126d130d7f71a85", size = 4401696, upload-time = "2026-04-08T01:57:31.029Z" }, + { url = "https://files.pythonhosted.org/packages/41/52/a8908dcb1a389a459a29008c29966c1d552588d4ae6d43f3a1a4512e0ebe/cryptography-46.0.7-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:a1529d614f44b863a7b480c6d000fe93b59acee9c82ffa027cfadc77521a9f5e", size = 4664256, upload-time = "2026-04-08T01:57:33.144Z" }, + { url = "https://files.pythonhosted.org/packages/4b/fa/f0ab06238e899cc3fb332623f337a7364f36f4bb3f2534c2bb95a35b132c/cryptography-46.0.7-cp38-abi3-win32.whl", hash = "sha256:f247c8c1a1fb45e12586afbb436ef21ff1e80670b2861a90353d9b025583d246", size = 3013001, upload-time = "2026-04-08T01:57:34.933Z" }, + { url = "https://files.pythonhosted.org/packages/d2/f1/00ce3bde3ca542d1acd8f8cfa38e446840945aa6363f9b74746394b14127/cryptography-46.0.7-cp38-abi3-win_amd64.whl", hash = "sha256:506c4ff91eff4f82bdac7633318a526b1d1309fc07ca76a3ad182cb5b686d6d3", size = 3472985, upload-time = "2026-04-08T01:57:36.714Z" }, +] + +[[package]] +name = "cuda-bindings" +version = "13.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-pathfinder", marker = "sys_platform != 'darwin'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/52/c8/b2589d68acf7e3d63e2be330b84bc25712e97ed799affbca7edd7eae25d6/cuda_bindings-13.2.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e865447abfb83d6a98ad5130ed3c70b1fc295ae3eeee39fd07b4ddb0671b6788", size = 5722404, upload-time = "2026-03-11T00:12:44.041Z" }, + { url = "https://files.pythonhosted.org/packages/1f/92/f899f7bbb5617bb65ec52a6eac1e9a1447a86b916c4194f8a5001b8cde0c/cuda_bindings-13.2.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:46d8776a55d6d5da9dd6e9858fba2efcda2abe6743871dee47dd06eb8cb6d955", size = 6320619, upload-time = "2026-03-11T00:12:45.939Z" }, + { url = "https://files.pythonhosted.org/packages/df/93/eef988860a3ca985f82c4f3174fc0cdd94e07331ba9a92e8e064c260337f/cuda_bindings-13.2.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6629ca2df6f795b784752409bcaedbd22a7a651b74b56a165ebc0c9dcbd504d0", size = 5614610, upload-time = "2026-03-11T00:12:50.337Z" }, + { url = "https://files.pythonhosted.org/packages/18/23/6db3aba46864aee357ab2415135b3fe3da7e9f1fa0221fa2a86a5968099c/cuda_bindings-13.2.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7dca0da053d3b4cc4869eff49c61c03f3c5dbaa0bcd712317a358d5b8f3f385d", size = 6149914, upload-time = "2026-03-11T00:12:52.374Z" }, +] + +[[package]] +name = "cuda-pathfinder" +version = "1.5.3" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d3/d6/ac63065d33dd700fee7ebd7d287332401b54e31b9346e142f871e1f0b116/cuda_pathfinder-1.5.3-py3-none-any.whl", hash = "sha256:dff021123aedbb4117cc7ec81717bbfe198fb4e8b5f1ee57e0e084fec5c8577d", size = 49991, upload-time = "2026-04-14T20:09:27.037Z" }, +] + +[[package]] +name = "cuda-toolkit" +version = "13.0.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/57/b2/453099f5f3b698d7d0eab38916aac44c7f76229f451709e2eb9db6615dcd/cuda_toolkit-13.0.2-py2.py3-none-any.whl", hash = "sha256:b198824cf2f54003f50d64ada3a0f184b42ca0846c1c94192fa269ecd97a66eb", size = 2364, upload-time = "2025-12-19T23:24:07.328Z" }, +] + +[package.optional-dependencies] +cublas = [ + { name = "nvidia-cublas", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +cudart = [ + { name = "nvidia-cuda-runtime", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +cufft = [ + { name = "nvidia-cufft", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +cufile = [ + { name = "nvidia-cufile", marker = "sys_platform == 'linux'" }, +] +cupti = [ + { name = "nvidia-cuda-cupti", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +curand = [ + { name = "nvidia-curand", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +cusolver = [ + { name = "nvidia-cusolver", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +cusparse = [ + { name = "nvidia-cusparse", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +nvjitlink = [ + { name = "nvidia-nvjitlink", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +nvrtc = [ + { name = "nvidia-cuda-nvrtc", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +nvtx = [ + { name = "nvidia-nvtx", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] + +[[package]] +name = "decorator" +version = "5.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/43/fa/6d96a0978d19e17b68d634497769987b16c8f4cd0a7a05048bec693caa6b/decorator-5.2.1.tar.gz", hash = "sha256:65f266143752f734b0a7cc83c46f4618af75b8c5911b00ccb61d0ac9b6da0360", size = 56711, upload-time = "2025-02-24T04:41:34.073Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190, upload-time = "2025-02-24T04:41:32.565Z" }, +] + +[[package]] +name = "deprecated" +version = "1.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wrapt" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/49/85/12f0a49a7c4ffb70572b6c2ef13c90c88fd190debda93b23f026b25f9634/deprecated-1.3.1.tar.gz", hash = "sha256:b1b50e0ff0c1fddaa5708a2c6b0a6588bb09b892825ab2b214ac9ea9d92a5223", size = 2932523, upload-time = "2025-10-30T08:19:02.757Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/84/d0/205d54408c08b13550c733c4b85429e7ead111c7f0014309637425520a9a/deprecated-1.3.1-py2.py3-none-any.whl", hash = "sha256:597bfef186b6f60181535a29fbe44865ce137a5079f295b479886c82729d5f3f", size = 11298, upload-time = "2025-10-30T08:19:00.758Z" }, +] + +[[package]] +name = "docstring-parser" +version = "0.18.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/4d/f332313098c1de1b2d2ff91cf2674415cc7cddab2ca1b01ae29774bd5fdf/docstring_parser-0.18.0.tar.gz", hash = "sha256:292510982205c12b1248696f44959db3cdd1740237a968ea1e2e7a900eeb2015", size = 29341, upload-time = "2026-04-14T04:09:19.867Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/5f/ed01f9a3cdffbd5a008556fc7b2a08ddb1cc6ace7effa7340604b1d16699/docstring_parser-0.18.0-py3-none-any.whl", hash = "sha256:b3fcbed555c47d8479be0796ef7e19c2670d428d72e96da63f3a40122860374b", size = 22484, upload-time = "2026-04-14T04:09:18.638Z" }, +] + +[[package]] +name = "embreex" +version = "2.17.7.post7" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", marker = "platform_machine != 's390x'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/01/ed299e72f62731be03c383ee4e13191d28afd7e73a8a30012e5b12c61916/embreex-2.17.7.post7-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b10cddba4cc294122e5a76c6df8b042fcb291b352220ecc72b2d9559fc45c48b", size = 10738000, upload-time = "2025-10-22T20:07:29.599Z" }, + { url = "https://files.pythonhosted.org/packages/75/c9/8f7ee075d31f51c9e2c50a1a380cf4cbb5c1d56c55a3d26dc098cf9bef5a/embreex-2.17.7.post7-cp312-cp312-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e069c7140ab8d4a3712fcc902dafda31e36c5aaa5fdc3fe5ba0f0e613cbe1876", size = 17235388, upload-time = "2025-10-22T20:07:32.121Z" }, + { url = "https://files.pythonhosted.org/packages/bb/47/483c92bd50a334b81d94f7a6a224080be0e20705a1d76bfb75935ee1ce8c/embreex-2.17.7.post7-cp312-cp312-win_amd64.whl", hash = "sha256:239e34700d821a55365565bcb2c9275237bc32984d0e518e04e106278978911e", size = 13466148, upload-time = "2025-10-22T20:07:34.334Z" }, + { url = "https://files.pythonhosted.org/packages/99/1a/e6d80ab0104ba947e1e5a0c5251021a4bf374cfe3b0d88e2be62b150170d/embreex-2.17.7.post7-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:377888d4f61b6d523e91b835aea3bd74193bd2b5fee9212adfe5fa0d5b80671b", size = 10734165, upload-time = "2025-10-22T20:07:36.707Z" }, + { url = "https://files.pythonhosted.org/packages/59/c7/410b62268fa98d9f0f2f5ab14b6e718e308a5f30dc4b9c2dcb0223d98ac8/embreex-2.17.7.post7-cp313-cp313-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c947025954a43a17a51d17be5ddbcfd1881ec5d0ea443bdf0ac4ac67e8e15e15", size = 17225042, upload-time = "2025-10-22T20:07:38.833Z" }, + { url = "https://files.pythonhosted.org/packages/c8/17/cf4bfc4f30cdaa9b64670e223d55949aae4f3d0005a8ba6559c7327477ac/embreex-2.17.7.post7-cp313-cp313-win_amd64.whl", hash = "sha256:c0f125c0517a4200a3f5d1979920704b924ec1ff6d6d0d826238216e40bf65fe", size = 13463316, upload-time = "2025-10-22T20:07:42.26Z" }, +] + +[[package]] +name = "etils" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/26/ce/6e067242fde898841922ac6fc82b0bb2fe35c38e995880bdffdfbe30182a/etils-1.14.0.tar.gz", hash = "sha256:8136e7f4c4173cd0af0ca5481c4475152f0b8686192951eefa60ee8711e1ede4", size = 108127, upload-time = "2026-03-04T17:41:36.291Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl", hash = "sha256:b5df7341f54dbe1405a4450b2741207b4a8c279780402b45f87202b94dfc52b4", size = 172934, upload-time = "2026-03-04T17:41:35.01Z" }, +] + +[package.optional-dependencies] +epath = [ + { name = "fsspec" }, + { name = "typing-extensions" }, + { name = "zipp" }, +] + +[[package]] +name = "evdev" +version = "1.9.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a5/f5/397b61091120a9ca5001041dd7bf76c385b3bfd67a0e5bcb74b852bd22a4/evdev-1.9.3.tar.gz", hash = "sha256:2c140e01ac8437758fa23fe5c871397412461f42d421aa20241dc8fe8cfccbc9", size = 32723, upload-time = "2026-02-05T21:54:24.987Z" } + +[[package]] +name = "fabric" +version = "3.2.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "decorator" }, + { name = "deprecated" }, + { name = "invoke" }, + { name = "paramiko" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e3/7e/29cd6237c3b7ce79c3ca945eb99ab5affd101db54b2f7a78dde0cfa19fd4/fabric-3.2.3.tar.gz", hash = "sha256:dcbd2c47ad87688facaef5cc11aab6d1ec9ed05645fed97a5de7204d5d17cc44", size = 183497, upload-time = "2026-04-06T00:00:11.481Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/37/f9/f8497ef8b873a8bb2a750ee2a6c5f0fc22258e1acb6245fd237042a6c279/fabric-3.2.3-py3-none-any.whl", hash = "sha256:ce61917f4f398018337ce279b357650a3a74baecf3fdd53a5839013944af965e", size = 59502, upload-time = "2026-04-06T00:00:10.176Z" }, +] + +[[package]] +name = "filelock" +version = "3.29.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b5/fe/997687a931ab51049acce6fa1f23e8f01216374ea81374ddee763c493db5/filelock-3.29.0.tar.gz", hash = "sha256:69974355e960702e789734cb4871f884ea6fe50bd8404051a3530bc07809cf90", size = 57571, upload-time = "2026-04-19T15:39:10.068Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/81/47/dd9a212ef6e343a6857485ffe25bba537304f1913bdbed446a23f7f592e1/filelock-3.29.0-py3-none-any.whl", hash = "sha256:96f5f6344709aa1572bbf631c640e4ebeeb519e08da902c39a001882f30ac258", size = 39812, upload-time = "2026-04-19T15:39:08.752Z" }, +] + +[[package]] +name = "flatbuffers" +version = "25.12.19" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/2d/d2a548598be01649e2d46231d151a6c56d10b964d94043a335ae56ea2d92/flatbuffers-25.12.19-py2.py3-none-any.whl", hash = "sha256:7634f50c427838bb021c2d66a3d1168e9d199b0607e6329399f04846d42e20b4", size = 26661, upload-time = "2025-12-19T23:16:13.622Z" }, +] + +[[package]] +name = "fsspec" +version = "2026.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e1/cf/b50ddf667c15276a9ab15a70ef5f257564de271957933ffea49d2cdbcdfb/fsspec-2026.3.0.tar.gz", hash = "sha256:1ee6a0e28677557f8c2f994e3eea77db6392b4de9cd1f5d7a9e87a0ae9d01b41", size = 313547, upload-time = "2026-03-27T19:11:14.892Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d5/1f/5f4a3cd9e4440e9d9bc78ad0a91a1c8d46b4d429d5239ebe6793c9fe5c41/fsspec-2026.3.0-py3-none-any.whl", hash = "sha256:d2ceafaad1b3457968ed14efa28798162f1638dbb5d2a6868a2db002a5ee39a4", size = 202595, upload-time = "2026-03-27T19:11:13.595Z" }, +] + +[[package]] +name = "gitdb" +version = "4.0.12" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "smmap" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/72/94/63b0fc47eb32792c7ba1fe1b694daec9a63620db1e313033d18140c2320a/gitdb-4.0.12.tar.gz", hash = "sha256:5ef71f855d191a3326fcfbc0d5da835f26b13fbcba60c32c21091c349ffdb571", size = 394684, upload-time = "2025-01-02T07:20:46.413Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/61/5c78b91c3143ed5c14207f463aecfc8f9dbb5092fb2869baf37c273b2705/gitdb-4.0.12-py3-none-any.whl", hash = "sha256:67073e15955400952c6565cc3e707c554a4eea2e428946f7a4c162fab9bd9bcf", size = 62794, upload-time = "2025-01-02T07:20:43.624Z" }, +] + +[[package]] +name = "gitpython" +version = "3.1.46" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "gitdb" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/b5/59d16470a1f0dfe8c793f9ef56fd3826093fc52b3bd96d6b9d6c26c7e27b/gitpython-3.1.46.tar.gz", hash = "sha256:400124c7d0ef4ea03f7310ac2fbf7151e09ff97f2a3288d64a440c584a29c37f", size = 215371, upload-time = "2026-01-01T15:37:32.073Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6a/09/e21df6aef1e1ffc0c816f0522ddc3f6dcded766c3261813131c78a704470/gitpython-3.1.46-py3-none-any.whl", hash = "sha256:79812ed143d9d25b6d176a10bb511de0f9c67b1fa641d82097b0ab90398a2058", size = 208620, upload-time = "2026-01-01T15:37:30.574Z" }, +] + +[[package]] +name = "glfw" +version = "2.10.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/96/72/642d4f12f61816ac96777f7360d413e3977a7dd08237d196f02da681b186/glfw-2.10.0.tar.gz", hash = "sha256:801e55d8581b34df9aa2cfea43feb06ff617576e2a8cc5dac23ee75b26d10abe", size = 31475, upload-time = "2025-09-12T08:54:38.871Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3d/1f/a9ce08b1173b0ab625ee92f0c47a5278b3e76fd367699880d8ee7d56c338/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-macosx_10_6_intel.whl", hash = "sha256:5f365a8c94bcea71ec91327e7c16e7cf739128479a18b8c1241b004b40acc412", size = 105329, upload-time = "2025-09-12T08:54:27.938Z" }, + { url = "https://files.pythonhosted.org/packages/7c/96/5a2220abcbd027eebcf8bedd28207a2de168899e51be13ba01ebdd4147a1/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-macosx_11_0_arm64.whl", hash = "sha256:5328db1a92d07abd988730517ec02aa8390d3e6ef7ce98c8b57ecba2f43a39ba", size = 102179, upload-time = "2025-09-12T08:54:29.163Z" }, + { url = "https://files.pythonhosted.org/packages/9d/41/a5bd1d9e1808f400102bd7d328c4ac17b65fb2fc8014014ec6f23d02f662/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-manylinux2014_aarch64.whl", hash = "sha256:312c4c1dd5509613ed6bc1e95a8dbb75a36b6dcc4120f50dc3892b40172e9053", size = 230039, upload-time = "2025-09-12T08:54:30.201Z" }, + { url = "https://files.pythonhosted.org/packages/80/aa/3b503c448609dee6cb4e7138b4109338f0e65b97be107ab85562269d378d/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-manylinux2014_x86_64.whl", hash = "sha256:59c53387dc08c62e8bed86bbe3a8d53ab1b27161281ffa0e7f27b64284e2627c", size = 241984, upload-time = "2025-09-12T08:54:31.347Z" }, + { url = "https://files.pythonhosted.org/packages/ac/2d/bfe39a42cad8e80b02bf5f7cae19ba67832c1810bbd3624a8e83153d74a4/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-manylinux_2_28_aarch64.whl", hash = "sha256:c6f292fdaf3f9a99e598ede6582d21c523a6f51f8f5e66213849101a6bcdc699", size = 231052, upload-time = "2025-09-12T08:54:32.859Z" }, + { url = "https://files.pythonhosted.org/packages/f7/02/6e639e90f181dc9127046e00d0528f9f7ad12d428972e3a5378b9aefdb0b/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-manylinux_2_28_x86_64.whl", hash = "sha256:7916034efa867927892635733a3b6af8cd95ceb10566fd7f1e0d2763c2ee8b12", size = 243525, upload-time = "2025-09-12T08:54:34.006Z" }, + { url = "https://files.pythonhosted.org/packages/84/06/cb588ca65561defe0fc48d1df4c2ac12569b81231ae4f2b52ab37007d0bd/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-win32.whl", hash = "sha256:6c9549da71b93e367b4d71438798daae1da2592039fd14204a80a1a2348ae127", size = 552685, upload-time = "2025-09-12T08:54:35.723Z" }, + { url = "https://files.pythonhosted.org/packages/86/27/00c9c96af18ac0a5eac2ff61cbe306551a2d770d7173f396d0792ee1a59e/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.p39.p310.p311.p312.p313-none-win_amd64.whl", hash = "sha256:6292d5d6634d668cd23d337e6089491d3945a9aa4ac6e1667b0003520d7caa51", size = 559466, upload-time = "2025-09-12T08:54:37.661Z" }, + { url = "https://files.pythonhosted.org/packages/b3/87/de0b33f6f00687499ca1371f22aa73396341b85bf88f1a284f9da8842493/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-macosx_10_6_intel.whl", hash = "sha256:2aab89d2d9535635ba011fc7303390685169a1aa6731ad580d08d043524b8899", size = 105326, upload-time = "2026-01-28T05:57:56.083Z" }, + { url = "https://files.pythonhosted.org/packages/b6/a6/6ea2f73ad4474896d9e38b3ffbe6ffd5a802c738392269e99e8c6621a461/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-macosx_11_0_arm64.whl", hash = "sha256:23936202a107039b5372f0b88ae1d11080746aa1c78910a45d4a0c4cf408cfaa", size = 102180, upload-time = "2026-01-28T05:57:57.787Z" }, + { url = "https://files.pythonhosted.org/packages/58/19/d81b19e8261b9cb51b81d1402167791fef81088dfe91f0c4e9d136fdc5ca/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-manylinux2014_aarch64.whl", hash = "sha256:7be06d0838f61df67bd54cb6266a6193d54083acb3624ff3c3812a6358406fa4", size = 230038, upload-time = "2026-01-28T05:57:59.105Z" }, + { url = "https://files.pythonhosted.org/packages/e2/fa/b035636cd82198b97b51a93efe9cfc4343d6b15cefbd336a3f2be871d848/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-manylinux2014_x86_64.whl", hash = "sha256:91d36b3582a766512eff8e3b5dcc2d3ffcbf10b7cf448551085a08a10f1b8244", size = 241983, upload-time = "2026-01-28T05:58:00.352Z" }, + { url = "https://files.pythonhosted.org/packages/ff/b4/f7b6cc022dd7c68b6c702d19da5d591f978f89c958b9bd3090615db0c739/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-manylinux_2_28_aarch64.whl", hash = "sha256:27c9e9a2d5e1dc3c9e3996171d844d9df9a5a101e797cb94cce217b7afcf8fd9", size = 231053, upload-time = "2026-01-28T05:58:01.683Z" }, + { url = "https://files.pythonhosted.org/packages/5a/3f/efeb7c6801c46e11bd666a5180f0d615f74f72264212f74f39586c6fda9d/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-manylinux_2_28_x86_64.whl", hash = "sha256:ce6724bb7cb3d0543dcba17206dce909f94176e68220b8eafee72e9f92bcf542", size = 243522, upload-time = "2026-01-28T05:58:03.517Z" }, + { url = "https://files.pythonhosted.org/packages/cf/b9/b04c3aa0aad2870cfe799f32f8b59789c98e1816bbce9e83f4823c5b840b/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-win32.whl", hash = "sha256:fca724a21a372731edb290841edd28a9fb1ee490f833392752844ac807c0086a", size = 552682, upload-time = "2026-01-28T05:58:05.649Z" }, + { url = "https://files.pythonhosted.org/packages/bd/e1/6d6816b296a529ac9b897ad228b1e084eb1f92319e96371880eebdc874a6/glfw-2.10.0-py2.py27.py3.py30.py31.py32.py33.py34.py35.py36.py37.py38.py39.py310.py311.py312.py313.py314-none-win_amd64.whl", hash = "sha256:823c0bd7770977d4b10e0ed0aef2f3682276b7c88b8b65cfc540afce5951392f", size = 559464, upload-time = "2026-01-28T05:58:07.261Z" }, + { url = "https://files.pythonhosted.org/packages/8c/a8/d4dab8a58fc2e6981fc7a58c4e56ba9d777fb24931cec6a22152edbb3540/glfw-2.10.0-py2.py3-none-macosx_10_6_intel.whl", hash = "sha256:a0d1f29f206219cc291edfb6cace663a86da2470632551c998e3db82d48ea177", size = 105288, upload-time = "2026-03-10T17:21:19.929Z" }, + { url = "https://files.pythonhosted.org/packages/14/61/68d35e001872a7705112418da236fa2418d4f2e5419f8b2837f9b81bb3da/glfw-2.10.0-py2.py3-none-macosx_11_0_arm64.whl", hash = "sha256:d28d6f3ef217e64e35dc6fd0a7acb4cec9bfe7cd14dd9b35a7228a87002de154", size = 102139, upload-time = "2026-03-10T17:21:21.645Z" }, + { url = "https://files.pythonhosted.org/packages/4e/e1/ca5984081aaae07c9d371cb11dc4e4ff603510678ed9b73e58b6c351fe63/glfw-2.10.0-py2.py3-none-manylinux2014_aarch64.whl", hash = "sha256:f968b522bb6a0e04aaf4dcac30a476d7229308bb2bac406a60587debb5a61e29", size = 229998, upload-time = "2026-03-10T17:21:23.549Z" }, + { url = "https://files.pythonhosted.org/packages/fa/c4/82ac75fdcfba2896da7a573c0fc7f8ceb8f77ead6866d500d06c32f1c464/glfw-2.10.0-py2.py3-none-manylinux2014_x86_64.whl", hash = "sha256:68cf3752bdadb6f4bc0a876247c28c88c7251ac39f8af076ed938fdfd71e72dd", size = 241944, upload-time = "2026-03-10T17:21:26.102Z" }, + { url = "https://files.pythonhosted.org/packages/e3/96/9f691823cca5eb6a08f346bd0ff03b78032db9370b509a1e9c8976fb20a5/glfw-2.10.0-py2.py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:44d98de5dbf8f727e0cb29f9b29d29528ea7570f2e6f42f8430a69df05f12b48", size = 231009, upload-time = "2026-03-10T17:21:28.481Z" }, + { url = "https://files.pythonhosted.org/packages/3f/93/977b9e679e356871d428ae7a1139ec767dd5177bed58a6344b4d2199e00f/glfw-2.10.0-py2.py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:cca5158d62189e08792b1ae54f92307a282921a0e7783315b467e21b0a381c88", size = 243480, upload-time = "2026-03-10T17:21:30.538Z" }, + { url = "https://files.pythonhosted.org/packages/f9/bd/cea9569c8f2188b0a104472951420434a3e1f5cf26f5836ef9d7227a1a30/glfw-2.10.0-py2.py3-none-win32.whl", hash = "sha256:5e024509989740e8e7b86cc4aab508195495f79879072b0e1f68bd036a2916ad", size = 552641, upload-time = "2026-03-10T17:21:32.653Z" }, + { url = "https://files.pythonhosted.org/packages/cc/9b/4366ad3e1c0688146c70aa6143584d6a8d88583b9390f106250e25a3d5cd/glfw-2.10.0-py2.py3-none-win_amd64.whl", hash = "sha256:7f787ee8645781f10e8800438ce4357ab38c573ffb191aba380c1e72eba6311c", size = 559423, upload-time = "2026-03-10T17:21:34.766Z" }, +] + +[[package]] +name = "grpcio" +version = "1.80.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b7/48/af6173dbca4454f4637a4678b67f52ca7e0c1ed7d5894d89d434fecede05/grpcio-1.80.0.tar.gz", hash = "sha256:29aca15edd0688c22ba01d7cc01cb000d72b2033f4a3c72a81a19b56fd143257", size = 12978905, upload-time = "2026-03-30T08:49:10.502Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5c/e8/a2b749265eb3415abc94f2e619bbd9e9707bebdda787e61c593004ec927a/grpcio-1.80.0-cp312-cp312-linux_armv7l.whl", hash = "sha256:c624cc9f1008361014378c9d776de7182b11fe8b2e5a81bc69f23a295f2a1ad0", size = 6015616, upload-time = "2026-03-30T08:47:13.428Z" }, + { url = "https://files.pythonhosted.org/packages/3e/97/b1282161a15d699d1e90c360df18d19165a045ce1c343c7f313f5e8a0b77/grpcio-1.80.0-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:f49eddcac43c3bf350c0385366a58f36bed8cc2c0ec35ef7b74b49e56552c0c2", size = 12014204, upload-time = "2026-03-30T08:47:15.873Z" }, + { url = "https://files.pythonhosted.org/packages/6e/5e/d319c6e997b50c155ac5a8cb12f5173d5b42677510e886d250d50264949d/grpcio-1.80.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d334591df610ab94714048e0d5b4f3dd5ad1bee74dfec11eee344220077a79de", size = 6563866, upload-time = "2026-03-30T08:47:18.588Z" }, + { url = "https://files.pythonhosted.org/packages/ae/f6/fdd975a2cb4d78eb67769a7b3b3830970bfa2e919f1decf724ae4445f42c/grpcio-1.80.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:0cb517eb1d0d0aaf1d87af7cc5b801d686557c1d88b2619f5e31fab3c2315921", size = 7273060, upload-time = "2026-03-30T08:47:21.113Z" }, + { url = "https://files.pythonhosted.org/packages/db/f0/a3deb5feba60d9538a962913e37bd2e69a195f1c3376a3dd44fe0427e996/grpcio-1.80.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4e78c4ac0d97dc2e569b2f4bcbbb447491167cb358d1a389fc4af71ab6f70411", size = 6782121, upload-time = "2026-03-30T08:47:23.827Z" }, + { url = "https://files.pythonhosted.org/packages/ca/84/36c6dcfddc093e108141f757c407902a05085e0c328007cb090d56646cdf/grpcio-1.80.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2ed770b4c06984f3b47eb0517b1c69ad0b84ef3f40128f51448433be904634cd", size = 7383811, upload-time = "2026-03-30T08:47:26.517Z" }, + { url = "https://files.pythonhosted.org/packages/7c/ef/f3a77e3dc5b471a0ec86c564c98d6adfa3510d38f8ee99010410858d591e/grpcio-1.80.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:256507e2f524092f1473071a05e65a5b10d84b82e3ff24c5b571513cfaa61e2f", size = 8393860, upload-time = "2026-03-30T08:47:29.439Z" }, + { url = "https://files.pythonhosted.org/packages/9b/8d/9d4d27ed7f33d109c50d6b5ce578a9914aa68edab75d65869a17e630a8d1/grpcio-1.80.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:9a6284a5d907c37db53350645567c522be314bac859a64a7a5ca63b77bb7958f", size = 7830132, upload-time = "2026-03-30T08:47:33.254Z" }, + { url = "https://files.pythonhosted.org/packages/14/e4/9990b41c6d7a44e1e9dee8ac11d7a9802ba1378b40d77468a7761d1ad288/grpcio-1.80.0-cp312-cp312-win32.whl", hash = "sha256:c71309cfce2f22be26aa4a847357c502db6c621f1a49825ae98aa0907595b193", size = 4140904, upload-time = "2026-03-30T08:47:35.319Z" }, + { url = "https://files.pythonhosted.org/packages/2f/2c/296f6138caca1f4b92a31ace4ae1b87dab692fc16a7a3417af3bb3c805bf/grpcio-1.80.0-cp312-cp312-win_amd64.whl", hash = "sha256:9fe648599c0e37594c4809d81a9e77bd138cc82eb8baa71b6a86af65426723ff", size = 4880944, upload-time = "2026-03-30T08:47:37.831Z" }, + { url = "https://files.pythonhosted.org/packages/2f/3a/7c3c25789e3f069e581dc342e03613c5b1cb012c4e8c7d9d5cf960a75856/grpcio-1.80.0-cp313-cp313-linux_armv7l.whl", hash = "sha256:e9e408fc016dffd20661f0126c53d8a31c2821b5c13c5d67a0f5ed5de93319ad", size = 6017243, upload-time = "2026-03-30T08:47:40.075Z" }, + { url = "https://files.pythonhosted.org/packages/04/19/21a9806eb8240e174fd1ab0cd5b9aa948bb0e05c2f2f55f9d5d7405e6d08/grpcio-1.80.0-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:92d787312e613754d4d8b9ca6d3297e69994a7912a32fa38c4c4e01c272974b0", size = 12010840, upload-time = "2026-03-30T08:47:43.11Z" }, + { url = "https://files.pythonhosted.org/packages/18/3a/23347d35f76f639e807fb7a36fad3068aed100996849a33809591f26eca6/grpcio-1.80.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8ac393b58aa16991a2f1144ec578084d544038c12242da3a215966b512904d0f", size = 6567644, upload-time = "2026-03-30T08:47:46.806Z" }, + { url = "https://files.pythonhosted.org/packages/ff/40/96e07ecb604a6a67ae6ab151e3e35b132875d98bc68ec65f3e5ab3e781d7/grpcio-1.80.0-cp313-cp313-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:68e5851ac4b9afe07e7f84483803ad167852570d65326b34d54ca560bfa53fb6", size = 7277830, upload-time = "2026-03-30T08:47:49.643Z" }, + { url = "https://files.pythonhosted.org/packages/9b/e2/da1506ecea1f34a5e365964644b35edef53803052b763ca214ba3870c856/grpcio-1.80.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:873ff5d17d68992ef6605330127425d2fc4e77e612fa3c3e0ed4e668685e3140", size = 6783216, upload-time = "2026-03-30T08:47:52.817Z" }, + { url = "https://files.pythonhosted.org/packages/44/83/3b20ff58d0c3b7f6caaa3af9a4174d4023701df40a3f39f7f1c8e7c48f9d/grpcio-1.80.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:2bea16af2750fd0a899bf1abd9022244418b55d1f37da2202249ba4ba673838d", size = 7385866, upload-time = "2026-03-30T08:47:55.687Z" }, + { url = "https://files.pythonhosted.org/packages/47/45/55c507599c5520416de5eefecc927d6a0d7af55e91cfffb2e410607e5744/grpcio-1.80.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:ba0db34f7e1d803a878284cd70e4c63cb6ae2510ba51937bf8f45ba997cefcf7", size = 8391602, upload-time = "2026-03-30T08:47:58.303Z" }, + { url = "https://files.pythonhosted.org/packages/10/bb/dd06f4c24c01db9cf11341b547d0a016b2c90ed7dbbb086a5710df7dd1d7/grpcio-1.80.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8eb613f02d34721f1acf3626dfdb3545bd3c8505b0e52bf8b5710a28d02e8aa7", size = 7826752, upload-time = "2026-03-30T08:48:01.311Z" }, + { url = "https://files.pythonhosted.org/packages/f9/1e/9d67992ba23371fd63d4527096eb8c6b76d74d52b500df992a3343fd7251/grpcio-1.80.0-cp313-cp313-win32.whl", hash = "sha256:93b6f823810720912fd131f561f91f5fed0fda372b6b7028a2681b8194d5d294", size = 4142310, upload-time = "2026-03-30T08:48:04.594Z" }, + { url = "https://files.pythonhosted.org/packages/cf/e6/283326a27da9e2c3038bc93eeea36fb118ce0b2d03922a9cda6688f53c5b/grpcio-1.80.0-cp313-cp313-win_amd64.whl", hash = "sha256:e172cf795a3ba5246d3529e4d34c53db70e888fa582a8ffebd2e6e48bc0cba50", size = 4882833, upload-time = "2026-03-30T08:48:07.363Z" }, +] + +[[package]] +name = "h11" +version = "0.16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/01/ee/02a2c011bdab74c6fb3c75474d40b3052059d95df7e73351460c8588d963/h11-0.16.0.tar.gz", hash = "sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1", size = 101250, upload-time = "2025-04-24T03:35:25.427Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" }, +] + +[[package]] +name = "httpcore" +version = "1.0.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "h11" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/94/82699a10bca87a5556c9c59b5963f2d039dbd239f25bc2a63907a05a14cb/httpcore-1.0.9.tar.gz", hash = "sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8", size = 85484, upload-time = "2025-04-24T22:06:22.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" }, +] + +[[package]] +name = "httpx" +version = "0.28.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "certifi" }, + { name = "httpcore" }, + { name = "idna" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b1/df/48c586a5fe32a0f01324ee087459e112ebb7224f646c0b5023f5e79e9956/httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc", size = 141406, upload-time = "2024-12-06T15:37:23.222Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" }, +] + +[[package]] +name = "idna" +version = "3.12" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/22/12/2948fbe5513d062169bd91f7d7b1cd97bc8894f32946b71fa39f6e63ca0c/idna-3.12.tar.gz", hash = "sha256:724e9952cc9e2bd7550ea784adb098d837ab5267ef67a1ab9cf7846bdbdd8254", size = 194350, upload-time = "2026-04-21T13:32:48.916Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/53/b2/acc33950394b3becb2b664741a0c0889c7ef9f9ffbfa8d47eddb53a50abd/idna-3.12-py3-none-any.whl", hash = "sha256:60ffaa1858fac94c9c124728c24fcde8160f3fb4a7f79aa8cdd33a9d1af60a67", size = 68634, upload-time = "2026-04-21T13:32:47.403Z" }, +] + +[[package]] +name = "imageio" +version = "2.37.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "pillow" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b1/84/93bcd1300216ea50811cee96873b84a1bebf8d0489ffaf7f2a3756bab866/imageio-2.37.3.tar.gz", hash = "sha256:bbb37efbfc4c400fcd534b367b91fcd66d5da639aaa138034431a1c5e0a41451", size = 389673, upload-time = "2026-03-09T11:31:12.573Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/fa/391e437a34e55095173dca5f24070d89cbc233ff85bf1c29c93248c6588d/imageio-2.37.3-py3-none-any.whl", hash = "sha256:46f5bb8522cd421c0f5ae104d8268f569d856b29eb1a13b92829d1970f32c9f0", size = 317646, upload-time = "2026-03-09T11:31:10.771Z" }, +] + +[[package]] +name = "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 = "importlib-metadata" +version = "9.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "zipp" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a9/01/15bb152d77b21318514a96f43af312635eb2500c96b55398d020c93d86ea/importlib_metadata-9.0.0.tar.gz", hash = "sha256:a4f57ab599e6a2e3016d7595cfd72eb4661a5106e787a95bcc90c7105b831efc", size = 56405, upload-time = "2026-03-20T06:42:56.999Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/38/3d/2d244233ac4f76e38533cfcb2991c9eb4c7bf688ae0a036d30725b8faafe/importlib_metadata-9.0.0-py3-none-any.whl", hash = "sha256:2d21d1cc5a017bd0559e36150c21c830ab1dc304dedd1b7ea85d20f45ef3edd7", size = 27789, upload-time = "2026-03-20T06:42:55.665Z" }, +] + +[[package]] +name = "invoke" +version = "2.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/de/bd/b461d3424a24c80490313fd77feeb666ca4f6a28c7e72713e3d9095719b4/invoke-2.2.1.tar.gz", hash = "sha256:515bf49b4a48932b79b024590348da22f39c4942dff991ad1fb8b8baea1be707", size = 304762, upload-time = "2025-10-11T00:36:35.172Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/32/4b/b99e37f88336009971405cbb7630610322ed6fbfa31e1d7ab3fbf3049a2d/invoke-2.2.1-py3-none-any.whl", hash = "sha256:2413bc441b376e5cd3f55bb5d364f973ad8bdd7bf87e53c79de3c11bf3feecc8", size = 160287, upload-time = "2025-10-11T00:36:33.703Z" }, +] + +[[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 = "jsonschema" +version = "4.26.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, + { name = "jsonschema-specifications" }, + { name = "referencing" }, + { name = "rpds-py" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b3/fc/e067678238fa451312d4c62bf6e6cf5ec56375422aee02f9cb5f909b3047/jsonschema-4.26.0.tar.gz", hash = "sha256:0c26707e2efad8aa1bfc5b7ce170f3fccc2e4918ff85989ba9ffa9facb2be326", size = 366583, upload-time = "2026-01-07T13:41:07.246Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/69/90/f63fb5873511e014207a475e2bb4e8b2e570d655b00ac19a9a0ca0a385ee/jsonschema-4.26.0-py3-none-any.whl", hash = "sha256:d489f15263b8d200f8387e64b4c3a75f06629559fb73deb8fdfb525f2dab50ce", size = 90630, upload-time = "2026-01-07T13:41:05.306Z" }, +] + +[[package]] +name = "jsonschema-specifications" +version = "2025.9.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "referencing" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/19/74/a633ee74eb36c44aa6d1095e7cc5569bebf04342ee146178e2d36600708b/jsonschema_specifications-2025.9.1.tar.gz", hash = "sha256:b540987f239e745613c7a9176f3edb72b832a4ac465cf02712288397832b5e8d", size = 32855, upload-time = "2025-09-08T01:34:59.186Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/41/45/1a4ed80516f02155c51f51e8cedb3c1902296743db0bbc66608a0db2814f/jsonschema_specifications-2025.9.1-py3-none-any.whl", hash = "sha256:98802fee3a11ee76ecaca44429fda8a41bff98b00a0f2838151b113f210cc6fe", size = 18437, upload-time = "2025-09-08T01:34:57.871Z" }, +] + +[[package]] +name = "lxml" +version = "6.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/28/30/9abc9e34c657c33834eaf6cd02124c61bdf5944d802aa48e69be8da3585d/lxml-6.1.0.tar.gz", hash = "sha256:bfd57d8008c4965709a919c3e9a98f76c2c7cb319086b3d26858250620023b13", size = 4197006, upload-time = "2026-04-18T04:32:51.613Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d2/d4/9326838b59dc36dfae42eec9656b97520f9997eee1de47b8316aaeed169c/lxml-6.1.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:d2f17a16cd8751e8eb233a7e41aecdf8e511712e00088bf9be455f604cd0d28d", size = 8570663, upload-time = "2026-04-18T04:27:48.253Z" }, + { url = "https://files.pythonhosted.org/packages/d8/a4/053745ce1f8303ccbb788b86c0db3a91b973675cefc42566a188637b7c40/lxml-6.1.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:f0cea5b1d3e6e77d71bd2b9972eb2446221a69dc52bb0b9c3c6f6e5700592d93", size = 4624024, upload-time = "2026-04-18T04:27:52.594Z" }, + { url = "https://files.pythonhosted.org/packages/90/97/a517944b20f8fd0932ad2109482bee4e29fe721416387a363306667941f6/lxml-6.1.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fc46da94826188ed45cb53bd8e3fc076ae22675aea2087843d4735627f867c6d", size = 4930895, upload-time = "2026-04-18T04:32:56.29Z" }, + { url = "https://files.pythonhosted.org/packages/94/7c/e08a970727d556caa040a44773c7b7e3ad0f0d73dedc863543e9a8b931f2/lxml-6.1.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9147d8e386ec3b82c3b15d88927f734f565b0aaadef7def562b853adca45784a", size = 5093820, upload-time = "2026-04-18T04:32:58.94Z" }, + { url = "https://files.pythonhosted.org/packages/88/ee/2a5c2aa2c32016a226ca25d3e1056a8102ea6e1fe308bf50213586635400/lxml-6.1.0-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5715e0e28736a070f3f34a7ccc09e2fdcba0e3060abbcf61a1a5718ff6d6b105", size = 5005790, upload-time = "2026-04-18T04:33:01.272Z" }, + { url = "https://files.pythonhosted.org/packages/e3/38/a0db9be8f38ad6043ab9429487c128dd1d30f07956ef43040402f8da49e8/lxml-6.1.0-cp312-cp312-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:4937460dc5df0cdd2f06a86c285c28afda06aefa3af949f9477d3e8df430c485", size = 5630827, upload-time = "2026-04-18T04:33:04.036Z" }, + { url = "https://files.pythonhosted.org/packages/31/ba/3c13d3fc24b7cacf675f808a3a1baabf43a30d0cd24c98f94548e9aa58eb/lxml-6.1.0-cp312-cp312-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bc783ee3147e60a25aa0445ea82b3e8aabb83b240f2b95d32cb75587ff781814", size = 5240445, upload-time = "2026-04-18T04:33:06.87Z" }, + { url = "https://files.pythonhosted.org/packages/55/ba/eeef4ccba09b2212fe239f46c1692a98db1878e0872ae320756488878a94/lxml-6.1.0-cp312-cp312-manylinux_2_28_i686.whl", hash = "sha256:40d9189f80075f2e1f88db21ef815a2b17b28adf8e50aaf5c789bfe737027f32", size = 5350121, upload-time = "2026-04-18T04:33:09.365Z" }, + { url = "https://files.pythonhosted.org/packages/7e/01/1da87c7b587c38d0cbe77a01aae3b9c1c49ed47d76918ef3db8fc151b1ca/lxml-6.1.0-cp312-cp312-manylinux_2_31_armv7l.whl", hash = "sha256:05b9b8787e35bec69e68daf4952b2e6dfcfb0db7ecf1a06f8cdfbbac4eb71aad", size = 4694949, upload-time = "2026-04-18T04:33:11.628Z" }, + { url = "https://files.pythonhosted.org/packages/a1/88/7db0fe66d5aaf128443ee1623dec3db1576f3e4c17751ec0ef5866468590/lxml-6.1.0-cp312-cp312-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:0f0f08beb0182e3e9a86fae124b3c47a7b41b7b69b225e1377db983802404e54", size = 5243901, upload-time = "2026-04-18T04:33:13.95Z" }, + { url = "https://files.pythonhosted.org/packages/00/a8/1346726af7d1f6fca1f11223ba34001462b0a3660416986d37641708d57c/lxml-6.1.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:73becf6d8c81d4c76b1014dbd3584cb26d904492dcf73ca85dc8bff08dcd6d2d", size = 5048054, upload-time = "2026-04-18T04:33:16.965Z" }, + { url = "https://files.pythonhosted.org/packages/2e/b7/85057012f035d1a0c87e02f8c723ca3c3e6e0728bcf4cb62080b21b1c1e3/lxml-6.1.0-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:1ae225f66e5938f4fa29d37e009a3bb3b13032ac57eb4eb42afa44f6e4054e69", size = 4777324, upload-time = "2026-04-18T04:33:19.832Z" }, + { url = "https://files.pythonhosted.org/packages/75/6c/ad2f94a91073ef570f33718040e8e160d5fb93331cf1ab3ca1323f939e2d/lxml-6.1.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:690022c7fae793b0489aa68a658822cea83e0d5933781811cabbf5ea3bcfe73d", size = 5645702, upload-time = "2026-04-18T04:33:22.436Z" }, + { url = "https://files.pythonhosted.org/packages/3b/89/0bb6c0bd549c19004c60eea9dc554dd78fd647b72314ef25d460e0d208c6/lxml-6.1.0-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:63aeafc26aac0be8aff14af7871249e87ea1319be92090bfd632ec68e03b16a5", size = 5232901, upload-time = "2026-04-18T04:33:26.21Z" }, + { url = "https://files.pythonhosted.org/packages/a1/d9/d609a11fb567da9399f525193e2b49847b5a409cdebe737f06a8b7126bdc/lxml-6.1.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:264c605ab9c0e4aa1a679636f4582c4d3313700009fac3ec9c3412ed0d8f3e1d", size = 5261333, upload-time = "2026-04-18T04:33:28.984Z" }, + { url = "https://files.pythonhosted.org/packages/a6/3a/ac3f99ec8ac93089e7dd556f279e0d14c24de0a74a507e143a2e4b496e7c/lxml-6.1.0-cp312-cp312-win32.whl", hash = "sha256:56971379bc5ee8037c5a0f09fa88f66cdb7d37c3e38af3e45cf539f41131ac1f", size = 3596289, upload-time = "2026-04-18T04:27:42.819Z" }, + { url = "https://files.pythonhosted.org/packages/f2/a7/0a915557538593cb1bbeedcd40e13c7a261822c26fecbbdb71dad0c2f540/lxml-6.1.0-cp312-cp312-win_amd64.whl", hash = "sha256:bba078de0031c219e5dd06cf3e6bf8fb8e6e64a77819b358f53bb132e3e03366", size = 3997059, upload-time = "2026-04-18T04:27:46.764Z" }, + { url = "https://files.pythonhosted.org/packages/92/96/a5dc078cf0126fbfbc35611d77ecd5da80054b5893e28fb213a5613b9e1d/lxml-6.1.0-cp312-cp312-win_arm64.whl", hash = "sha256:c3592631e652afa34999a088f98ba7dfc7d6aff0d535c410bea77a71743f3819", size = 3659552, upload-time = "2026-04-18T04:27:51.133Z" }, + { url = "https://files.pythonhosted.org/packages/08/03/69347590f1cf4a6d5a4944bb6099e6d37f334784f16062234e1f892fdb1d/lxml-6.1.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:a0092f2b107b69601adf562a57c956fbb596e05e3e6651cabd3054113b007e45", size = 8559689, upload-time = "2026-04-18T04:31:57.785Z" }, + { url = "https://files.pythonhosted.org/packages/3f/58/25e00bb40b185c974cfe156c110474d9a8a8390d5f7c92a4e328189bb60e/lxml-6.1.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:fc7140d7a7386e6b545d41b7358f4d02b656d4053f5fa6859f92f4b9c2572c4d", size = 4617892, upload-time = "2026-04-18T04:32:01.78Z" }, + { url = "https://files.pythonhosted.org/packages/f5/54/92ad98a94ac318dc4f97aaac22ff8d1b94212b2ae8af5b6e9b354bf825f7/lxml-6.1.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:419c58fc92cc3a2c3fa5f78c63dbf5da70c1fa9c1b25f25727ecee89a96c7de2", size = 4923489, upload-time = "2026-04-18T04:33:31.401Z" }, + { url = "https://files.pythonhosted.org/packages/15/3b/a20aecfab42bdf4f9b390590d345857ad3ffd7c51988d1c89c53a0c73faf/lxml-6.1.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:37fabd1452852636cf38ecdcc9dd5ca4bba7a35d6c53fa09725deeb894a87491", size = 5082162, upload-time = "2026-04-18T04:33:34.262Z" }, + { url = "https://files.pythonhosted.org/packages/45/26/2cdb3d281ac1bd175603e290cbe4bad6eff127c0f8de90bafd6f8548f0fd/lxml-6.1.0-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a2853c8b2170cc6cd54a6b4d50d2c1a8a7aeca201f23804b4898525c7a152cfc", size = 4993247, upload-time = "2026-04-18T04:33:36.674Z" }, + { url = "https://files.pythonhosted.org/packages/f6/05/d735aef963740022a08185c84821f689fc903acb3d50326e6b1e9886cc22/lxml-6.1.0-cp313-cp313-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:8e369cbd690e788c8d15e56222d91a09c6a417f49cbc543040cba0fe2e25a79e", size = 5613042, upload-time = "2026-04-18T04:33:39.205Z" }, + { url = "https://files.pythonhosted.org/packages/ee/b8/ead7c10efff731738c72e59ed6eb5791854879fbed7ae98781a12006263a/lxml-6.1.0-cp313-cp313-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e69aa6805905807186eb00e66c6d97a935c928275182eb02ee40ba00da9623b2", size = 5228304, upload-time = "2026-04-18T04:33:41.647Z" }, + { url = "https://files.pythonhosted.org/packages/6b/10/e9842d2ec322ea65f0a7270aa0315a53abed06058b88ef1b027f620e7a5f/lxml-6.1.0-cp313-cp313-manylinux_2_28_i686.whl", hash = "sha256:4bd1bdb8a9e0e2dd229de19b5f8aebac80e916921b4b2c6ef8a52bc131d0c1f9", size = 5341578, upload-time = "2026-04-18T04:33:44.596Z" }, + { url = "https://files.pythonhosted.org/packages/89/54/40d9403d7c2775fa7301d3ddd3464689bfe9ba71acc17dfff777071b4fdc/lxml-6.1.0-cp313-cp313-manylinux_2_31_armv7l.whl", hash = "sha256:cbd7b79cdcb4986ad78a2662625882747f09db5e4cd7b2ae178a88c9c51b3dfe", size = 4700209, upload-time = "2026-04-18T04:33:47.552Z" }, + { url = "https://files.pythonhosted.org/packages/85/b2/bbdcc2cf45dfc7dfffef4fd97e5c47b15919b6a365247d95d6f684ef5e82/lxml-6.1.0-cp313-cp313-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:43e4d297f11080ec9d64a4b1ad7ac02b4484c9f0e2179d9c4ef78e886e747b88", size = 5232365, upload-time = "2026-04-18T04:33:50.249Z" }, + { url = "https://files.pythonhosted.org/packages/48/5a/b06875665e53aaba7127611a7bed3b7b9658e20b22bc2dd217a0b7ab0091/lxml-6.1.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:cc16682cc987a3da00aa56a3aa3075b08edb10d9b1e476938cfdbee8f3b67181", size = 5043654, upload-time = "2026-04-18T04:33:52.71Z" }, + { url = "https://files.pythonhosted.org/packages/e9/9c/e71a069d09641c1a7abeb30e693f828c7c90a41cbe3d650b2d734d876f85/lxml-6.1.0-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:d6d8efe71429635f0559579092bb5e60560d7b9115ee38c4adbea35632e7fa24", size = 4769326, upload-time = "2026-04-18T04:33:55.244Z" }, + { url = "https://files.pythonhosted.org/packages/cc/06/7a9cd84b3d4ed79adf35f874750abb697dec0b4a81a836037b36e47c091a/lxml-6.1.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:7e39ab3a28af7784e206d8606ec0e4bcad0190f63a492bca95e94e5a4aef7f6e", size = 5635879, upload-time = "2026-04-18T04:33:58.509Z" }, + { url = "https://files.pythonhosted.org/packages/cc/f0/9d57916befc1e54c451712c7ee48e9e74e80ae4d03bdce49914e0aee42cd/lxml-6.1.0-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:9eb667bf50856c4a58145f8ca2d5e5be160191e79eb9e30855a476191b3c3495", size = 5224048, upload-time = "2026-04-18T04:34:00.943Z" }, + { url = "https://files.pythonhosted.org/packages/99/75/90c4eefda0c08c92221fe0753db2d6699a4c628f76ff4465ec20dea84cc1/lxml-6.1.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:7f4a77d6f7edf9230cee3e1f7f6764722a41604ee5681844f18db9a81ea0ec33", size = 5250241, upload-time = "2026-04-18T04:34:03.365Z" }, + { url = "https://files.pythonhosted.org/packages/5e/73/16596f7e4e38fa33084b9ccbccc22a15f82a290a055126f2c1541236d2ff/lxml-6.1.0-cp313-cp313-win32.whl", hash = "sha256:28902146ffbe5222df411c5d19e5352490122e14447e98cd118907ee3fd6ee62", size = 3596938, upload-time = "2026-04-18T04:31:56.206Z" }, + { url = "https://files.pythonhosted.org/packages/8e/63/981401c5680c1eb30893f00a19641ac80db5d1e7086c62cb4b13ed813038/lxml-6.1.0-cp313-cp313-win_amd64.whl", hash = "sha256:4a1503c56e4e2b38dc76f2f2da7bae69670c0f1933e27cfa34b2fa5876410b16", size = 3995728, upload-time = "2026-04-18T04:31:58.763Z" }, + { url = "https://files.pythonhosted.org/packages/e7/e8/c358a38ac3e541d16a1b527e4e9cb78c0419b0506a070ace11777e5e8404/lxml-6.1.0-cp313-cp313-win_arm64.whl", hash = "sha256:e0af85773850417d994d019741239b901b22c6680206f46a34766926e466141d", size = 3658372, upload-time = "2026-04-18T04:32:03.629Z" }, +] + +[[package]] +name = "manifold3d" +version = "3.4.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b0/fd/4dfc246e076e3912c45a821764f4de8b6c8117fa36fc67f8e44bf9dfe59b/manifold3d-3.4.1.tar.gz", hash = "sha256:b517927e2c15dc52169fff0cd12e1949eceb4ca49f3a5b8c0568b1116a561ab1", size = 269275, upload-time = "2026-03-24T06:22:40.062Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d9/59/def4c589dd55aa32026a720f8a31d71aa2162fef8e3963b6241a7945ef4e/manifold3d-3.4.1-cp312-cp312-macosx_10_14_universal2.whl", hash = "sha256:967c89daf24ec9ff863323d593cce98e4c130abbaaa9504df6789f9d8c780d0d", size = 1752517, upload-time = "2026-03-24T06:21:55.203Z" }, + { url = "https://files.pythonhosted.org/packages/b1/a9/377800999cc8421ce8bfa40787d09570bb635e0099f44959170fee751dd7/manifold3d-3.4.1-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:c29db9a1bda414ecaa56dd2cd274f06bbbe740e463133c5b69943d82c3dcfb96", size = 956343, upload-time = "2026-03-24T06:21:57.134Z" }, + { url = "https://files.pythonhosted.org/packages/6a/72/7f988a0deae9b3fbed3a6b2e9285e96fd9105e95f6755f5457e3a80e103e/manifold3d-3.4.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:2ae6855a6f652acd89e228f1981e5b710d4b10e06d7c06e5bada3b3fb31904a3", size = 840924, upload-time = "2026-03-24T06:21:58.462Z" }, + { url = "https://files.pythonhosted.org/packages/bf/46/787ad4b53a35ccf1d31fbd3d2ffe0653dab67057b1f561db51d2edc494ba/manifold3d-3.4.1-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a8068f85416034e290d23b424f2bea15d2f1da1c5ccb79b442bdb50ed4e1a4b6", size = 1253014, upload-time = "2026-03-24T06:21:59.734Z" }, + { url = "https://files.pythonhosted.org/packages/e2/45/29d2380ac477b11629a72483b21dd544861caccaedbc87043bf315a15a50/manifold3d-3.4.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fe9ff5ce3d949c72b21120318121eb926ddba4299eab0e8bab2c6784a9843ffb", size = 1355512, upload-time = "2026-03-24T06:22:01.518Z" }, + { url = "https://files.pythonhosted.org/packages/7f/7f/310688a725a5a23d00e9f29e614a2b7906b399df27731b1aa6e153e4f465/manifold3d-3.4.1-cp312-cp312-win_amd64.whl", hash = "sha256:1bd6fa1b20238603ec3df7f6060ddba1181cf9464104e82b746747b487d12092", size = 983293, upload-time = "2026-03-24T06:22:03.273Z" }, + { url = "https://files.pythonhosted.org/packages/e3/d0/b066b476242dddfad98db51425a28ac41ab008a4e7697f6d6bca21a24881/manifold3d-3.4.1-cp313-cp313-macosx_10_14_universal2.whl", hash = "sha256:210ec6918870611d9e3f888c00657aad842cfa89a7967e94546a568bf8dfc2f1", size = 1752343, upload-time = "2026-03-24T06:22:04.731Z" }, + { url = "https://files.pythonhosted.org/packages/11/2c/10b5cb142b00bb7b14bb5698c584ce7722c68c3ed58ae4173693a35d2108/manifold3d-3.4.1-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:d4e1dd76a3c5fe935bc14eb62f98ce2361e0e505fcfca08abb2f9d8a1d01e0db", size = 956241, upload-time = "2026-03-24T06:22:06.404Z" }, + { url = "https://files.pythonhosted.org/packages/05/a7/af84e5f6e6af2d07d800355345ffb303c4e8de96dbf3194633322f3d8335/manifold3d-3.4.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:7e1bf4857f64311fb5113ff286898c47efc0f199e4d860cfc663b5f69ce90ede", size = 840900, upload-time = "2026-03-24T06:22:07.974Z" }, + { url = "https://files.pythonhosted.org/packages/80/1c/f274d6e35652c3fb72b54c5fcafc5fc474e1a93d3fd17fb8df3c9c765873/manifold3d-3.4.1-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9b23435456d5ed64e48a34a281869c3b626da8ccd8872e54637b77f420716f9a", size = 1252521, upload-time = "2026-03-24T06:22:09.259Z" }, + { url = "https://files.pythonhosted.org/packages/97/90/82081bcbffc68e36f9f34c36f041d6e0176cbb462e9041683d82ff17b626/manifold3d-3.4.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6871adaf9e5081303d2c9446de5a76a9af84bcf938949fa198cc5f0ae9cad19f", size = 1355366, upload-time = "2026-03-24T06:22:10.892Z" }, + { url = "https://files.pythonhosted.org/packages/0d/db/26df1d96a2c61a4d79aeb0ca2f8bfbfd4af94fdb944469dda38ced240f2f/manifold3d-3.4.1-cp313-cp313-win_amd64.whl", hash = "sha256:0a93e8202cccea16c76a6c3a7d02300755cccea6536874ccfc160f8c4d8948c4", size = 983317, upload-time = "2026-03-24T06:22:12.257Z" }, +] + +[[package]] +name = "mapbox-earcut" +version = "2.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/bc/7b/bbf6b00488662be5d2eb7a188222c264b6f713bac10dc4a77bf37a4cb4b6/mapbox_earcut-2.0.0.tar.gz", hash = "sha256:81eab6b86cf99551deb698b98e3f7502c57900e5c479df15e1bdaf1a57f0f9d6", size = 39934, upload-time = "2025-11-16T18:41:27.251Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/93/846804029d955c3c841d8efff77c2b0e8d9aab057d3a077dc8e3f88b5ea4/mapbox_earcut-2.0.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:db55ce18e698bc9d90914ee7d4f8c3e4d23827456ece7c5d7a1ec91e90c7122b", size = 55623, upload-time = "2025-11-16T18:40:32.113Z" }, + { url = "https://files.pythonhosted.org/packages/d3/f6/cc9ece104bc3876b350dba6fef7f34fb7b20ecc028d2cdbdbecb436b1ed1/mapbox_earcut-2.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:01dd6099d16123baf582a11b2bd1d59ce848498cf0cdca3812fd1f8b20ff33b7", size = 52028, upload-time = "2025-11-16T18:40:33.516Z" }, + { url = "https://files.pythonhosted.org/packages/88/6e/230da4aabcc56c99e9bddb4c43ce7d4ba3609c0caf2d316fb26535d7c60c/mapbox_earcut-2.0.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2d5a098aae26a52282bc981a38e7bf6b889d2ea7442f2cd1903d2ba842f4ff07", size = 56351, upload-time = "2025-11-16T18:40:35.217Z" }, + { url = "https://files.pythonhosted.org/packages/1a/f7/5cdd3752526e91d91336c7263af7767b291d21e63c89d7190a60051f0f87/mapbox_earcut-2.0.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:de35f241d0b9110ad9260f295acedd9d7cc0d7acfe30d36b1b3ee8419c2caba1", size = 59209, upload-time = "2025-11-16T18:40:36.634Z" }, + { url = "https://files.pythonhosted.org/packages/7b/a2/b7781416cb93b37b95d0444e03f87184de8815e57ff202ce4105fa921325/mapbox_earcut-2.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:6cb63ab85e2e430c350f93e75c13f8b91cb8c8a045f3cd714c390b69a720368a", size = 152316, upload-time = "2025-11-16T18:40:38.147Z" }, + { url = "https://files.pythonhosted.org/packages/c1/74/396338e3d345e4e36fb23a0380921098b6a95ce7fb19c4777f4185a5974e/mapbox_earcut-2.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fb3c9f069fc3795306db87f8139f70c4f047532f897a3de05f54dc1faebc97f6", size = 157268, upload-time = "2025-11-16T18:40:39.753Z" }, + { url = "https://files.pythonhosted.org/packages/56/2c/66fd137ea86c508f6cd7247f7f6e2d1dabffc9f0e9ccf14c71406b197af1/mapbox_earcut-2.0.0-cp312-cp312-win32.whl", hash = "sha256:eb290e6676217707ed238dd55e07b0a8ca3ab928f6a27c4afefb2ff3af08d7cb", size = 51226, upload-time = "2025-11-16T18:40:41.018Z" }, + { url = "https://files.pythonhosted.org/packages/b8/84/7b78e37b0c2109243c0dad7d9ba9774b02fcee228bf61cf727a5aa1702e2/mapbox_earcut-2.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:5ef5b3319a43375272ad2cad9333ed16e569b5102e32a4241451358897e6f6ee", size = 56417, upload-time = "2025-11-16T18:40:42.173Z" }, + { url = "https://files.pythonhosted.org/packages/75/7f/cd7195aa27c1c8f2b9d38025a5a8663f32cd01c07b648a54b1308ab26c15/mapbox_earcut-2.0.0-cp312-cp312-win_arm64.whl", hash = "sha256:a4a3706feb5cc8c782d8f68bb0110c8d551304043f680a87a54b0651a2c208c3", size = 50111, upload-time = "2025-11-16T18:40:43.334Z" }, + { url = "https://files.pythonhosted.org/packages/8b/7c/c5dd5b255b9828ba5df729e62fdd470a322c938f07ef392ca03c0592bb3a/mapbox_earcut-2.0.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:582329a81bd36cf0f82e443c395bcb8cfdb10caddafec76acaebac7c20bf1c31", size = 55619, upload-time = "2025-11-16T18:40:44.44Z" }, + { url = "https://files.pythonhosted.org/packages/1a/3f/03f23eac9831e7d0d8da3d6993695a9a3724659c94e9997f6b7aaccc199d/mapbox_earcut-2.0.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d2ac5f610b3e44a3a0c4df06b5552d503b4f1c2c409eeca20dbe05112bd60955", size = 52023, upload-time = "2025-11-16T18:40:45.857Z" }, + { url = "https://files.pythonhosted.org/packages/39/f3/a92ccee494b3e437e4bd81ecd358e39d231dc90af010d6c43930506c10ad/mapbox_earcut-2.0.0-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:58cc88513b87734b243d86f0d3fb87e96e0a78d9abd8fd615c55f766dd63f949", size = 56357, upload-time = "2025-11-16T18:40:47.27Z" }, + { url = "https://files.pythonhosted.org/packages/03/30/e54ececd0403a5495c340b693075abec92a6d17dc44283b6cb059534f7ed/mapbox_earcut-2.0.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:40218d887798451932f3c335992834aa807c35cd497c6e0733470fdbd77f9521", size = 59215, upload-time = "2025-11-16T18:40:48.682Z" }, + { url = "https://files.pythonhosted.org/packages/6e/e1/8fbff13a074c1fbf702b30ce7ec4d878bc664d659c1c2b1697831f4ea3a8/mapbox_earcut-2.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:39fa5cfa0e855b028ec9b0200c88ebfa252448f343ce2f67b6fc07fe1f22a3ae", size = 152304, upload-time = "2025-11-16T18:40:49.85Z" }, + { url = "https://files.pythonhosted.org/packages/b9/d5/c757030b3cb3a9f2278ded6f7312d2b9d3761db6f3da8d395f7f7303dd66/mapbox_earcut-2.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:476b558473b8a43f238d46e819bc0f830c427842ec5feb19e23b4dcac8ad2455", size = 157270, upload-time = "2025-11-16T18:40:51.093Z" }, + { url = "https://files.pythonhosted.org/packages/96/63/589c6decb1f032d8811f1066da552f0a718830f592e6d6539fa4c3c766b8/mapbox_earcut-2.0.0-cp313-cp313-win32.whl", hash = "sha256:8c2d125c182acbc490b39503c0dec4f937bae180d0849a26bcea0ee4a76024bd", size = 51207, upload-time = "2025-11-16T18:40:52.285Z" }, + { url = "https://files.pythonhosted.org/packages/76/75/a79a6020c46d4f07731e88ec5cc9324f6b43343aba835def1dc0bf59fecf/mapbox_earcut-2.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:e049e6a37c228d7a9cb2f54ae405aa21d35c5175d849530fb32064ddb38ad5ab", size = 56416, upload-time = "2025-11-16T18:40:53.474Z" }, + { url = "https://files.pythonhosted.org/packages/ce/5f/83e878c2b3e9e6db1f60b598a2cc5ed4c2b5bc8d281575c964869414a159/mapbox_earcut-2.0.0-cp313-cp313-win_arm64.whl", hash = "sha256:8a833d73d63d4b6291bbd8b4d2f551e87f663282cdc547ecbbd9b423849ee996", size = 50103, upload-time = "2025-11-16T18:40:54.954Z" }, +] + +[[package]] +name = "markdown" +version = "3.10.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2b/f4/69fa6ed85ae003c2378ffa8f6d2e3234662abd02c10d216c0ba96081a238/markdown-3.10.2.tar.gz", hash = "sha256:994d51325d25ad8aa7ce4ebaec003febcce822c3f8c911e3b17c52f7f589f950", size = 368805, upload-time = "2026-02-09T14:57:26.942Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/de/1f/77fa3081e4f66ca3576c896ae5d31c3002ac6607f9747d2e3aa49227e464/markdown-3.10.2-py3-none-any.whl", hash = "sha256:e91464b71ae3ee7afd3017d9f358ef0baf158fd9a298db92f1d4761133824c36", size = 108180, upload-time = "2026-02-09T14:57:25.787Z" }, +] + +[[package]] +name = "markdown-it-py" +version = "4.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mdurl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5b/f5/4ec618ed16cc4f8fb3b701563655a69816155e79e24a17b651541804721d/markdown_it_py-4.0.0.tar.gz", hash = "sha256:cb0a2b4aa34f932c007117b194e945bd74e0ec24133ceb5bac59009cda1cb9f3", size = 73070, upload-time = "2025-08-11T12:57:52.854Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl", hash = "sha256:87327c59b172c5011896038353a81343b6754500a08cd7a4973bb48c6d578147", size = 87321, upload-time = "2025-08-11T12:57:51.923Z" }, +] + +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e", size = 11615, upload-time = "2025-09-27T18:36:30.854Z" }, + { url = "https://files.pythonhosted.org/packages/9a/81/7e4e08678a1f98521201c3079f77db69fb552acd56067661f8c2f534a718/markupsafe-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce", size = 12020, upload-time = "2025-09-27T18:36:31.971Z" }, + { url = "https://files.pythonhosted.org/packages/1e/2c/799f4742efc39633a1b54a92eec4082e4f815314869865d876824c257c1e/markupsafe-3.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d", size = 24332, upload-time = "2025-09-27T18:36:32.813Z" }, + { url = "https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" }, + { url = "https://files.pythonhosted.org/packages/2c/54/887f3092a85238093a0b2154bd629c89444f395618842e8b0c41783898ea/markupsafe-3.0.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a", size = 21962, upload-time = "2025-09-27T18:36:35.099Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2f/336b8c7b6f4a4d95e91119dc8521402461b74a485558d8f238a68312f11c/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b", size = 23760, upload-time = "2025-09-27T18:36:36.001Z" }, + { url = "https://files.pythonhosted.org/packages/32/43/67935f2b7e4982ffb50a4d169b724d74b62a3964bc1a9a527f5ac4f1ee2b/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f", size = 21529, upload-time = "2025-09-27T18:36:36.906Z" }, + { url = "https://files.pythonhosted.org/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" }, + { url = "https://files.pythonhosted.org/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d", size = 14540, upload-time = "2025-09-27T18:36:38.761Z" }, + { url = "https://files.pythonhosted.org/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c", size = 15105, upload-time = "2025-09-27T18:36:39.701Z" }, + { url = "https://files.pythonhosted.org/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, + { url = "https://files.pythonhosted.org/packages/38/2f/907b9c7bbba283e68f20259574b13d005c121a0fa4c175f9bed27c4597ff/markupsafe-3.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e1cf1972137e83c5d4c136c43ced9ac51d0e124706ee1c8aa8532c1287fa8795", size = 11622, upload-time = "2025-09-27T18:36:41.777Z" }, + { url = "https://files.pythonhosted.org/packages/9c/d9/5f7756922cdd676869eca1c4e3c0cd0df60ed30199ffd775e319089cb3ed/markupsafe-3.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:116bb52f642a37c115f517494ea5feb03889e04df47eeff5b130b1808ce7c219", size = 12029, upload-time = "2025-09-27T18:36:43.257Z" }, + { url = "https://files.pythonhosted.org/packages/00/07/575a68c754943058c78f30db02ee03a64b3c638586fba6a6dd56830b30a3/markupsafe-3.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:133a43e73a802c5562be9bbcd03d090aa5a1fe899db609c29e8c8d815c5f6de6", size = 24374, upload-time = "2025-09-27T18:36:44.508Z" }, + { url = "https://files.pythonhosted.org/packages/a9/21/9b05698b46f218fc0e118e1f8168395c65c8a2c750ae2bab54fc4bd4e0e8/markupsafe-3.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ccfcd093f13f0f0b7fdd0f198b90053bf7b2f02a3927a30e63f3ccc9df56b676", size = 22980, upload-time = "2025-09-27T18:36:45.385Z" }, + { url = "https://files.pythonhosted.org/packages/7f/71/544260864f893f18b6827315b988c146b559391e6e7e8f7252839b1b846a/markupsafe-3.0.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:509fa21c6deb7a7a273d629cf5ec029bc209d1a51178615ddf718f5918992ab9", size = 21990, upload-time = "2025-09-27T18:36:46.916Z" }, + { url = "https://files.pythonhosted.org/packages/c2/28/b50fc2f74d1ad761af2f5dcce7492648b983d00a65b8c0e0cb457c82ebbe/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a4afe79fb3de0b7097d81da19090f4df4f8d3a2b3adaa8764138aac2e44f3af1", size = 23784, upload-time = "2025-09-27T18:36:47.884Z" }, + { url = "https://files.pythonhosted.org/packages/ed/76/104b2aa106a208da8b17a2fb72e033a5a9d7073c68f7e508b94916ed47a9/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:795e7751525cae078558e679d646ae45574b47ed6e7771863fcc079a6171a0fc", size = 21588, upload-time = "2025-09-27T18:36:48.82Z" }, + { url = "https://files.pythonhosted.org/packages/b5/99/16a5eb2d140087ebd97180d95249b00a03aa87e29cc224056274f2e45fd6/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8485f406a96febb5140bfeca44a73e3ce5116b2501ac54fe953e488fb1d03b12", size = 23041, upload-time = "2025-09-27T18:36:49.797Z" }, + { url = "https://files.pythonhosted.org/packages/19/bc/e7140ed90c5d61d77cea142eed9f9c303f4c4806f60a1044c13e3f1471d0/markupsafe-3.0.3-cp313-cp313-win32.whl", hash = "sha256:bdd37121970bfd8be76c5fb069c7751683bdf373db1ed6c010162b2a130248ed", size = 14543, upload-time = "2025-09-27T18:36:51.584Z" }, + { url = "https://files.pythonhosted.org/packages/05/73/c4abe620b841b6b791f2edc248f556900667a5a1cf023a6646967ae98335/markupsafe-3.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:9a1abfdc021a164803f4d485104931fb8f8c1efd55bc6b748d2f5774e78b62c5", size = 15113, upload-time = "2025-09-27T18:36:52.537Z" }, + { url = "https://files.pythonhosted.org/packages/f0/3a/fa34a0f7cfef23cf9500d68cb7c32dd64ffd58a12b09225fb03dd37d5b80/markupsafe-3.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:7e68f88e5b8799aa49c85cd116c932a1ac15caaa3f5db09087854d218359e485", size = 13911, upload-time = "2025-09-27T18:36:53.513Z" }, + { url = "https://files.pythonhosted.org/packages/e4/d7/e05cd7efe43a88a17a37b3ae96e79a19e846f3f456fe79c57ca61356ef01/markupsafe-3.0.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:218551f6df4868a8d527e3062d0fb968682fe92054e89978594c28e642c43a73", size = 11658, upload-time = "2025-09-27T18:36:54.819Z" }, + { url = "https://files.pythonhosted.org/packages/99/9e/e412117548182ce2148bdeacdda3bb494260c0b0184360fe0d56389b523b/markupsafe-3.0.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:3524b778fe5cfb3452a09d31e7b5adefeea8c5be1d43c4f810ba09f2ceb29d37", size = 12066, upload-time = "2025-09-27T18:36:55.714Z" }, + { url = "https://files.pythonhosted.org/packages/bc/e6/fa0ffcda717ef64a5108eaa7b4f5ed28d56122c9a6d70ab8b72f9f715c80/markupsafe-3.0.3-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4e885a3d1efa2eadc93c894a21770e4bc67899e3543680313b09f139e149ab19", size = 25639, upload-time = "2025-09-27T18:36:56.908Z" }, + { url = "https://files.pythonhosted.org/packages/96/ec/2102e881fe9d25fc16cb4b25d5f5cde50970967ffa5dddafdb771237062d/markupsafe-3.0.3-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8709b08f4a89aa7586de0aadc8da56180242ee0ada3999749b183aa23df95025", size = 23569, upload-time = "2025-09-27T18:36:57.913Z" }, + { url = "https://files.pythonhosted.org/packages/4b/30/6f2fce1f1f205fc9323255b216ca8a235b15860c34b6798f810f05828e32/markupsafe-3.0.3-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b8512a91625c9b3da6f127803b166b629725e68af71f8184ae7e7d54686a56d6", size = 23284, upload-time = "2025-09-27T18:36:58.833Z" }, + { url = "https://files.pythonhosted.org/packages/58/47/4a0ccea4ab9f5dcb6f79c0236d954acb382202721e704223a8aafa38b5c8/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9b79b7a16f7fedff2495d684f2b59b0457c3b493778c9eed31111be64d58279f", size = 24801, upload-time = "2025-09-27T18:36:59.739Z" }, + { url = "https://files.pythonhosted.org/packages/6a/70/3780e9b72180b6fecb83a4814d84c3bf4b4ae4bf0b19c27196104149734c/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:12c63dfb4a98206f045aa9563db46507995f7ef6d83b2f68eda65c307c6829eb", size = 22769, upload-time = "2025-09-27T18:37:00.719Z" }, + { url = "https://files.pythonhosted.org/packages/98/c5/c03c7f4125180fc215220c035beac6b9cb684bc7a067c84fc69414d315f5/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8f71bc33915be5186016f675cd83a1e08523649b0e33efdb898db577ef5bb009", size = 23642, upload-time = "2025-09-27T18:37:01.673Z" }, + { url = "https://files.pythonhosted.org/packages/80/d6/2d1b89f6ca4bff1036499b1e29a1d02d282259f3681540e16563f27ebc23/markupsafe-3.0.3-cp313-cp313t-win32.whl", hash = "sha256:69c0b73548bc525c8cb9a251cddf1931d1db4d2258e9599c28c07ef3580ef354", size = 14612, upload-time = "2025-09-27T18:37:02.639Z" }, + { url = "https://files.pythonhosted.org/packages/2b/98/e48a4bfba0a0ffcf9925fe2d69240bfaa19c6f7507b8cd09c70684a53c1e/markupsafe-3.0.3-cp313-cp313t-win_amd64.whl", hash = "sha256:1b4b79e8ebf6b55351f0d91fe80f893b4743f104bff22e90697db1590e47a218", size = 15200, upload-time = "2025-09-27T18:37:03.582Z" }, + { url = "https://files.pythonhosted.org/packages/0e/72/e3cc540f351f316e9ed0f092757459afbc595824ca724cbc5a5d4263713f/markupsafe-3.0.3-cp313-cp313t-win_arm64.whl", hash = "sha256:ad2cf8aa28b8c020ab2fc8287b0f823d0a7d8630784c31e9ee5edea20f406287", size = 13973, upload-time = "2025-09-27T18:37:04.929Z" }, +] + +[[package]] +name = "mdurl" +version = "0.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, +] + +[[package]] +name = "mjlab" +version = "1.0.0" +source = { directory = "/opt/vendor/mjlab" } +dependencies = [ + { name = "moviepy" }, + { name = "mujoco" }, + { name = "mujoco-warp" }, + { name = "onnxscript" }, + { name = "prettytable" }, + { name = "rsl-rl-lib" }, + { name = "tensorboard" }, + { name = "tensordict" }, + { name = "torch" }, + { name = "torchrunx" }, + { name = "tqdm" }, + { name = "trimesh" }, + { name = "tyro" }, + { name = "viser" }, + { name = "wandb" }, + { name = "warp-lang" }, +] + +[package.metadata] +requires-dist = [ + { name = "autodocsumm", marker = "extra == 'docs'" }, + { name = "moviepy" }, + { name = "mujoco", specifier = ">=3.4.0", index = "https://py.mujoco.org/" }, + { name = "mujoco-warp", directory = "/opt/vendor/mujoco_warp" }, + { name = "myst-parser", marker = "extra == 'docs'", specifier = ">=4.0.1" }, + { name = "onnxscript", specifier = ">=0.5.4" }, + { name = "prettytable" }, + { name = "rsl-rl-lib", specifier = "==3.1.0" }, + { name = "sphinx", marker = "extra == 'docs'", specifier = ">=8.1.3" }, + { name = "sphinx-autodoc-typehints", marker = "extra == 'docs'", specifier = ">=3.0.1" }, + { name = "sphinx-book-theme", marker = "extra == 'docs'", specifier = ">=1.1.4" }, + { name = "sphinx-copybutton", marker = "extra == 'docs'", specifier = ">=0.5.2" }, + { name = "sphinx-design", marker = "extra == 'docs'", specifier = ">=0.6.1" }, + { name = "sphinx-icon", marker = "extra == 'docs'" }, + { name = "sphinx-multiversion", marker = "extra == 'docs'" }, + { name = "sphinx-tabs", marker = "extra == 'docs'" }, + { name = "sphinxcontrib-bibtex", marker = "extra == 'docs'" }, + { name = "sphinxemoji", marker = "extra == 'docs'" }, + { name = "tensorboard", specifier = ">=2.20.0" }, + { name = "tensordict" }, + { name = "torch", specifier = ">=2.7.0" }, + { name = "torch", marker = "sys_platform == 'darwin' and extra == 'cu128'", specifier = ">=2.7.0" }, + { name = "torch", marker = "sys_platform != 'darwin' and extra == 'cu128'", specifier = ">=2.7.0", index = "https://download.pytorch.org/whl/cu128", conflict = { package = "mjlab", extra = "cu128" } }, + { name = "torchrunx", specifier = ">=0.3.4" }, + { name = "tqdm" }, + { name = "trimesh", specifier = ">=4.8.3" }, + { name = "tyro", specifier = ">=1.0.1" }, + { name = "viser", specifier = ">=1.0.21" }, + { name = "wandb", specifier = ">=0.22.3" }, + { name = "warp-lang", marker = "sys_platform != 'darwin'", specifier = "==1.13.0.dev20260415", index = "https://pypi.nvidia.com/" }, + { name = "warp-lang", marker = "sys_platform == 'darwin'", specifier = "==1.13.0.dev20260415" }, +] +provides-extras = ["cu128", "docs"] + +[package.metadata.requires-dev] +dev = [ + { name = "ipdb", specifier = ">=0.13.13" }, + { name = "pre-commit", specifier = ">=4.3.0" }, + { name = "pyright", specifier = ">=1.1.408" }, + { name = "pytest", specifier = ">=9.0.2" }, + { name = "ruff", specifier = ">=0.14.14" }, + { name = "ty", specifier = ">=0.0.14" }, +] + +[[package]] +name = "mjlab-husky" +version = "0.1.0" +source = { editable = "." } +dependencies = [ + { name = "mjlab" }, + { name = "onnxruntime" }, + { name = "pybullet" }, + { name = "pynput" }, + { name = "rerun-sdk" }, + { name = "rsl-rl-lib" }, +] + +[package.metadata] +requires-dist = [ + { name = "mjlab", directory = "/opt/vendor/mjlab" }, + { name = "onnxruntime" }, + { name = "pybullet" }, + { name = "pynput" }, + { name = "rerun-sdk" }, + { name = "rsl-rl-lib", directory = "rsl_rl" }, +] + +[[package]] +name = "ml-dtypes" +version = "0.5.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz", hash = "sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453", size = 692314, upload-time = "2025-11-17T22:32:31.031Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a8/b8/3c70881695e056f8a32f8b941126cf78775d9a4d7feba8abcb52cb7b04f2/ml_dtypes-0.5.4-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:a174837a64f5b16cab6f368171a1a03a27936b31699d167684073ff1c4237dac", size = 676927, upload-time = "2025-11-17T22:31:48.182Z" }, + { url = "https://files.pythonhosted.org/packages/54/0f/428ef6881782e5ebb7eca459689448c0394fa0a80bea3aa9262cba5445ea/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a7f7c643e8b1320fd958bf098aa7ecf70623a42ec5154e3be3be673f4c34d900", size = 5028464, upload-time = "2025-11-17T22:31:50.135Z" }, + { url = "https://files.pythonhosted.org/packages/3a/cb/28ce52eb94390dda42599c98ea0204d74799e4d8047a0eb559b6fd648056/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ad459e99793fa6e13bd5b7e6792c8f9190b4e5a1b45c63aba14a4d0a7f1d5ff", size = 5009002, upload-time = "2025-11-17T22:31:52.001Z" }, + { url = "https://files.pythonhosted.org/packages/f5/f0/0cfadd537c5470378b1b32bd859cf2824972174b51b873c9d95cfd7475a5/ml_dtypes-0.5.4-cp312-cp312-win_amd64.whl", hash = "sha256:c1a953995cccb9e25a4ae19e34316671e4e2edaebe4cf538229b1fc7109087b7", size = 212222, upload-time = "2025-11-17T22:31:53.742Z" }, + { url = "https://files.pythonhosted.org/packages/16/2e/9acc86985bfad8f2c2d30291b27cd2bb4c74cea08695bd540906ed744249/ml_dtypes-0.5.4-cp312-cp312-win_arm64.whl", hash = "sha256:9bad06436568442575beb2d03389aa7456c690a5b05892c471215bfd8cf39460", size = 160793, upload-time = "2025-11-17T22:31:55.358Z" }, + { url = "https://files.pythonhosted.org/packages/d9/a1/4008f14bbc616cfb1ac5b39ea485f9c63031c4634ab3f4cf72e7541f816a/ml_dtypes-0.5.4-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8c760d85a2f82e2bed75867079188c9d18dae2ee77c25a54d60e9cc79be1bc48", size = 676888, upload-time = "2025-11-17T22:31:56.907Z" }, + { url = "https://files.pythonhosted.org/packages/d3/b7/dff378afc2b0d5a7d6cd9d3209b60474d9819d1189d347521e1688a60a53/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ce756d3a10d0c4067172804c9cc276ba9cc0ff47af9078ad439b075d1abdc29b", size = 5036993, upload-time = "2025-11-17T22:31:58.497Z" }, + { url = "https://files.pythonhosted.org/packages/eb/33/40cd74219417e78b97c47802037cf2d87b91973e18bb968a7da48a96ea44/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:533ce891ba774eabf607172254f2e7260ba5f57bdd64030c9a4fcfbd99815d0d", size = 5010956, upload-time = "2025-11-17T22:31:59.931Z" }, + { url = "https://files.pythonhosted.org/packages/e1/8b/200088c6859d8221454825959df35b5244fa9bdf263fd0249ac5fb75e281/ml_dtypes-0.5.4-cp313-cp313-win_amd64.whl", hash = "sha256:f21c9219ef48ca5ee78402d5cc831bd58ea27ce89beda894428bc67a52da5328", size = 212224, upload-time = "2025-11-17T22:32:01.349Z" }, + { url = "https://files.pythonhosted.org/packages/8f/75/dfc3775cb36367816e678f69a7843f6f03bd4e2bcd79941e01ea960a068e/ml_dtypes-0.5.4-cp313-cp313-win_arm64.whl", hash = "sha256:35f29491a3e478407f7047b8a4834e4640a77d2737e0b294d049746507af5175", size = 160798, upload-time = "2025-11-17T22:32:02.864Z" }, + { url = "https://files.pythonhosted.org/packages/4f/74/e9ddb35fd1dd43b1106c20ced3f53c2e8e7fc7598c15638e9f80677f81d4/ml_dtypes-0.5.4-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:304ad47faa395415b9ccbcc06a0350800bc50eda70f0e45326796e27c62f18b6", size = 702083, upload-time = "2025-11-17T22:32:04.08Z" }, + { url = "https://files.pythonhosted.org/packages/74/f5/667060b0aed1aa63166b22897fdf16dca9eb704e6b4bbf86848d5a181aa7/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6a0df4223b514d799b8a1629c65ddc351b3efa833ccf7f8ea0cf654a61d1e35d", size = 5354111, upload-time = "2025-11-17T22:32:05.546Z" }, + { url = "https://files.pythonhosted.org/packages/40/49/0f8c498a28c0efa5f5c95a9e374c83ec1385ca41d0e85e7cf40e5d519a21/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:531eff30e4d368cb6255bc2328d070e35836aa4f282a0fb5f3a0cd7260257298", size = 5366453, upload-time = "2025-11-17T22:32:07.115Z" }, + { url = "https://files.pythonhosted.org/packages/8c/27/12607423d0a9c6bbbcc780ad19f1f6baa2b68b18ce4bddcdc122c4c68dc9/ml_dtypes-0.5.4-cp313-cp313t-win_amd64.whl", hash = "sha256:cb73dccfc991691c444acc8c0012bee8f2470da826a92e3a20bb333b1a7894e6", size = 225612, upload-time = "2025-11-17T22:32:08.615Z" }, + { url = "https://files.pythonhosted.org/packages/e5/80/5a5929e92c72936d5b19872c5fb8fc09327c1da67b3b68c6a13139e77e20/ml_dtypes-0.5.4-cp313-cp313t-win_arm64.whl", hash = "sha256:3bbbe120b915090d9dd1375e4684dd17a20a2491ef25d640a908281da85e73f1", size = 164145, upload-time = "2025-11-17T22:32:09.782Z" }, +] + +[[package]] +name = "moviepy" +version = "2.2.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "decorator" }, + { name = "imageio" }, + { name = "imageio-ffmpeg" }, + { name = "numpy" }, + { name = "pillow" }, + { name = "proglog" }, + { name = "python-dotenv" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/de/61/15f9476e270f64c78a834e7459ca045d669f869cec24eed26807b8cd479d/moviepy-2.2.1.tar.gz", hash = "sha256:c80cb56815ece94e5e3e2d361aa40070eeb30a09d23a24c4e684d03e16deacb1", size = 58431438, upload-time = "2025-05-21T19:31:52.601Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9a/73/7d3b2010baa0b5eb1e4dfa9e4385e89b6716be76f2fa21a6c0fe34b68e5a/moviepy-2.2.1-py3-none-any.whl", hash = "sha256:6b56803fec2ac54b557404126ac1160e65448e03798fa282bd23e8fab3795060", size = 129871, upload-time = "2025-05-21T19:31:50.11Z" }, +] + +[[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 = "msgspec" +version = "0.21.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e3/60/f79b9b013a16fa3a58350c9295ddc6789f2e335f36ea61ed10a21b215364/msgspec-0.21.1.tar.gz", hash = "sha256:2313508e394b0d208f8f56892ca9b2799e2561329de9763b19619595a6c0f72c", size = 319193, upload-time = "2026-04-12T21:44:50.394Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6e/cf/317224852c00248c620a9bcf4b26e2e4ab8afd752f18d2a6ef73ebd423b6/msgspec-0.21.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d4248cf0b6129b7d230eacd493c17cc2d4f3989f3bb7f633a928a85b7dcfa251", size = 196188, upload-time = "2026-04-12T21:44:07.181Z" }, + { url = "https://files.pythonhosted.org/packages/6d/81/074612945c0666078f7366f40000013de9f6ba687491d450df699bceebc9/msgspec-0.21.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:5102c7e9b3acff82178449b85006d96310e690291bb1ea0142f1b24bcb8aabcb", size = 188473, upload-time = "2026-04-12T21:44:08.736Z" }, + { url = "https://files.pythonhosted.org/packages/8a/37/655101799590bcc5fddb2bd3fe0e6194e816c2d1da7c361725f5eb89a910/msgspec-0.21.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:846758412e9518252b2ac9bffd6f0e54d9ff614f5f9488df7749f81ff5c80920", size = 218871, upload-time = "2026-04-12T21:44:09.917Z" }, + { url = "https://files.pythonhosted.org/packages/b5/d1/d4cd9fe89c7d400d7a18f86ccc94daa3f0927f53558846fcb60791dce5d6/msgspec-0.21.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:21995e74b5c598c2e004110ad66ec7f1b8c20bf2bcf3b2de8fd9a3094422d3ff", size = 225025, upload-time = "2026-04-12T21:44:11.191Z" }, + { url = "https://files.pythonhosted.org/packages/24/bf/e20549e602b9edccadeeff98760345a416f9cce846a657e8b18e3396b212/msgspec-0.21.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:6129f0cca52992e898fd5344187f7c8127b63d810b2fd73e36fca73b4c6475ee", size = 222672, upload-time = "2026-04-12T21:44:12.481Z" }, + { url = "https://files.pythonhosted.org/packages/b4/68/04d7a8f0f786545cf9b8c280c57aa6befb5977af6e884b8b54191cbe44b3/msgspec-0.21.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:ef3ec2296248d1f8b9231acb051b6d471dfde8f21819e86c9adaaa9f42918521", size = 227303, upload-time = "2026-04-12T21:44:13.709Z" }, + { url = "https://files.pythonhosted.org/packages/cc/4d/619866af2840875be408047bf9e70ceafbae6ab50660de7134ed1b25eb86/msgspec-0.21.1-cp312-cp312-win_amd64.whl", hash = "sha256:d4ab834a054c6f0cbeef6df9e7e1b33d5f1bc7b86dea1d2fd7cad003873e783d", size = 190017, upload-time = "2026-04-12T21:44:14.977Z" }, + { url = "https://files.pythonhosted.org/packages/5e/2e/a8f9eca8fd00e097d7a9e99ba8a4685db994494448e3d4f0b7f6e9a3c0f7/msgspec-0.21.1-cp312-cp312-win_arm64.whl", hash = "sha256:628aaa35c74950a8c59da330d7e98917e1c7188f983745782027748ee4ca573e", size = 175345, upload-time = "2026-04-12T21:44:16.431Z" }, + { url = "https://files.pythonhosted.org/packages/7e/74/f11ede02839b19ff459f88e3145df5d711626ca84da4e23520cebf819367/msgspec-0.21.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:764173717a01743f007e9f74520ed281f24672c604514f7d76c1c3a10e8edb66", size = 196176, upload-time = "2026-04-12T21:44:17.613Z" }, + { url = "https://files.pythonhosted.org/packages/bb/40/4476c1bd341418a046c4955aff632ec769315d1e3cb94e6acf86d461f9ed/msgspec-0.21.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:344c7cd0eaed1fb81d7959f99100ef71ec9b536881a376f11b9a6c4803365697", size = 188524, upload-time = "2026-04-12T21:44:18.815Z" }, + { url = "https://files.pythonhosted.org/packages/ca/d9/9e9d7d7e5061b47540d03d640fab9b3965ba7ae49c1b2154861c8f007518/msgspec-0.21.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:48943e278b3854c2f89f955ddc6f9f430d3f0784b16e47d10604ee0463cd21f5", size = 218880, upload-time = "2026-04-12T21:44:20.028Z" }, + { url = "https://files.pythonhosted.org/packages/74/66/2bb344f34abb4b57e60c7c9c761994e0417b9718ec1460bf00c296f2a7ea/msgspec-0.21.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a9aa659ebb0101b1cbc31461212b87e341d961f0ab0772aaf068a99e001ec4aa", size = 225050, upload-time = "2026-04-12T21:44:21.577Z" }, + { url = "https://files.pythonhosted.org/packages/1a/84/7c1e412f76092277bf760cef12b7979d03314d259ab5b5cafde5d0c1722d/msgspec-0.21.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f7b27d1a8ead2b6f5b0c4f2d07b8be1ccfcc041c8a0e704781edebe3ae13c484", size = 222713, upload-time = "2026-04-12T21:44:22.83Z" }, + { url = "https://files.pythonhosted.org/packages/4e/27/0bba04b2b4ef05f3d068429410bc71d2cea925f1596a8f41152cccd5edb8/msgspec-0.21.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:38fe93e86b61328fe544cb7fd871fad5a27c8734bfda90f65e5dbe288ae50f61", size = 227259, upload-time = "2026-04-12T21:44:24.11Z" }, + { url = "https://files.pythonhosted.org/packages/b0/2d/09574b0eea02fed2c2c1383dbaae2c7f79dc16dcd6487a886000afb5d7c4/msgspec-0.21.1-cp313-cp313-win_amd64.whl", hash = "sha256:8bc666331c35fcce05a7cd2d6221adbe0f6058f8e750711413d22793c080ac6a", size = 189857, upload-time = "2026-04-12T21:44:25.359Z" }, + { url = "https://files.pythonhosted.org/packages/46/34/105b1576ad182879914f0c821f17ee1d13abb165cb060448f96fe2aff078/msgspec-0.21.1-cp313-cp313-win_arm64.whl", hash = "sha256:42bb1241e0750c1a4346f2aa84db26c5ffd99a4eb3a954927d9f149ff2f42898", size = 175403, upload-time = "2026-04-12T21:44:26.608Z" }, +] + +[[package]] +name = "mujoco" +version = "3.7.1.dev903232388" +source = { registry = "https://py.mujoco.org/" } +dependencies = [ + { name = "absl-py" }, + { name = "etils", extra = ["epath"] }, + { name = "glfw" }, + { name = "numpy" }, + { name = "pyopengl" }, +] +wheels = [ + { url = "https://py.mujoco.org/mujoco/mujoco-3.7.1.dev903232388-cp312-cp312-macosx_11_0_universal2.whl" }, + { url = "https://py.mujoco.org/mujoco/mujoco-3.7.1.dev903232388-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl" }, + { url = "https://py.mujoco.org/mujoco/mujoco-3.7.1.dev903232388-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl" }, + { url = "https://py.mujoco.org/mujoco/mujoco-3.7.1.dev903232388-cp312-cp312-win_amd64.whl" }, + { url = "https://py.mujoco.org/mujoco/mujoco-3.7.1.dev903232388-cp313-cp313-macosx_11_0_universal2.whl" }, + { url = "https://py.mujoco.org/mujoco/mujoco-3.7.1.dev903232388-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl" }, + { url = "https://py.mujoco.org/mujoco/mujoco-3.7.1.dev903232388-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl" }, + { url = "https://py.mujoco.org/mujoco/mujoco-3.7.1.dev903232388-cp313-cp313-win_amd64.whl" }, +] + +[[package]] +name = "mujoco-warp" +version = "0.0.2" +source = { directory = "/opt/vendor/mujoco_warp" } +dependencies = [ + { name = "absl-py" }, + { name = "etils", extra = ["epath"] }, + { name = "mujoco" }, + { name = "numpy" }, + { name = "warp-lang" }, +] + +[package.metadata] +requires-dist = [ + { name = "absl-py" }, + { name = "etils", extras = ["epath"] }, + { name = "jax", marker = "extra == 'cpu'" }, + { name = "jax", extras = ["cuda12"], marker = "extra == 'cuda'" }, + { name = "lsprotocol", marker = "extra == 'dev'", specifier = ">=2023.0.1,<2024.0.0" }, + { name = "mujoco", specifier = ">=3.4.0", index = "https://py.mujoco.org/" }, + { name = "mujoco", marker = "extra == 'dev'", specifier = ">=3.3.7.dev0", index = "https://py.mujoco.org/" }, + { name = "numpy" }, + { name = "pre-commit", marker = "extra == 'dev'" }, + { name = "pygls", marker = "extra == 'dev'", specifier = ">=1.0.0,<2.0.0" }, + { name = "pytest", marker = "extra == 'dev'" }, + { name = "pytest-xdist", marker = "extra == 'dev'" }, + { name = "ruff", marker = "extra == 'dev'" }, + { name = "warp-lang", specifier = ">=1.11.0", index = "https://pypi.nvidia.com/" }, + { name = "warp-lang", marker = "extra == 'dev'", specifier = ">=1.11.0.dev0", index = "https://pypi.nvidia.com/" }, +] +provides-extras = ["dev", "cpu", "cuda"] + +[[package]] +name = "networkx" +version = "3.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, +] + +[[package]] +name = "numpy" +version = "2.4.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d7/9f/b8cef5bffa569759033adda9481211426f12f53299629b410340795c2514/numpy-2.4.4.tar.gz", hash = "sha256:2d390634c5182175533585cc89f3608a4682ccb173cc9bb940b2881c8d6f8fa0", size = 20731587, upload-time = "2026-03-29T13:22:01.298Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/28/05/32396bec30fb2263770ee910142f49c1476d08e8ad41abf8403806b520ce/numpy-2.4.4-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:15716cfef24d3a9762e3acdf87e27f58dc823d1348f765bbea6bef8c639bfa1b", size = 16689272, upload-time = "2026-03-29T13:18:49.223Z" }, + { url = "https://files.pythonhosted.org/packages/c5/f3/a983d28637bfcd763a9c7aafdb6d5c0ebf3d487d1e1459ffdb57e2f01117/numpy-2.4.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:23cbfd4c17357c81021f21540da84ee282b9c8fba38a03b7b9d09ba6b951421e", size = 14699573, upload-time = "2026-03-29T13:18:52.629Z" }, + { url = "https://files.pythonhosted.org/packages/9b/fd/e5ecca1e78c05106d98028114f5c00d3eddb41207686b2b7de3e477b0e22/numpy-2.4.4-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:8b3b60bb7cba2c8c81837661c488637eee696f59a877788a396d33150c35d842", size = 5204782, upload-time = "2026-03-29T13:18:55.579Z" }, + { url = "https://files.pythonhosted.org/packages/de/2f/702a4594413c1a8632092beae8aba00f1d67947389369b3777aed783fdca/numpy-2.4.4-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:e4a010c27ff6f210ff4c6ef34394cd61470d01014439b192ec22552ee867f2a8", size = 6552038, upload-time = "2026-03-29T13:18:57.769Z" }, + { url = "https://files.pythonhosted.org/packages/7f/37/eed308a8f56cba4d1fdf467a4fc67ef4ff4bf1c888f5fc980481890104b1/numpy-2.4.4-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f9e75681b59ddaa5e659898085ae0eaea229d054f2ac0c7e563a62205a700121", size = 15670666, upload-time = "2026-03-29T13:19:00.341Z" }, + { url = "https://files.pythonhosted.org/packages/0a/0d/0e3ecece05b7a7e87ab9fb587855548da437a061326fff64a223b6dcb78a/numpy-2.4.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:81f4a14bee47aec54f883e0cad2d73986640c1590eb9bfaaba7ad17394481e6e", size = 16645480, upload-time = "2026-03-29T13:19:03.63Z" }, + { url = "https://files.pythonhosted.org/packages/34/49/f2312c154b82a286758ee2f1743336d50651f8b5195db18cdb63675ff649/numpy-2.4.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:62d6b0f03b694173f9fcb1fb317f7222fd0b0b103e784c6549f5e53a27718c44", size = 17020036, upload-time = "2026-03-29T13:19:07.428Z" }, + { url = "https://files.pythonhosted.org/packages/7b/e9/736d17bd77f1b0ec4f9901aaec129c00d59f5d84d5e79bba540ef12c2330/numpy-2.4.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fbc356aae7adf9e6336d336b9c8111d390a05df88f1805573ebb0807bd06fd1d", size = 18368643, upload-time = "2026-03-29T13:19:10.775Z" }, + { url = "https://files.pythonhosted.org/packages/63/f6/d417977c5f519b17c8a5c3bc9e8304b0908b0e21136fe43bf628a1343914/numpy-2.4.4-cp312-cp312-win32.whl", hash = "sha256:0d35aea54ad1d420c812bfa0385c71cd7cc5bcf7c65fed95fc2cd02fe8c79827", size = 5961117, upload-time = "2026-03-29T13:19:13.464Z" }, + { url = "https://files.pythonhosted.org/packages/2d/5b/e1deebf88ff431b01b7406ca3583ab2bbb90972bbe1c568732e49c844f7e/numpy-2.4.4-cp312-cp312-win_amd64.whl", hash = "sha256:b5f0362dc928a6ecd9db58868fca5e48485205e3855957bdedea308f8672ea4a", size = 12320584, upload-time = "2026-03-29T13:19:16.155Z" }, + { url = "https://files.pythonhosted.org/packages/58/89/e4e856ac82a68c3ed64486a544977d0e7bdd18b8da75b78a577ca31c4395/numpy-2.4.4-cp312-cp312-win_arm64.whl", hash = "sha256:846300f379b5b12cc769334464656bc882e0735d27d9726568bc932fdc49d5ec", size = 10221450, upload-time = "2026-03-29T13:19:18.994Z" }, + { url = "https://files.pythonhosted.org/packages/14/1d/d0a583ce4fefcc3308806a749a536c201ed6b5ad6e1322e227ee4848979d/numpy-2.4.4-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:08f2e31ed5e6f04b118e49821397f12767934cfdd12a1ce86a058f91e004ee50", size = 16684933, upload-time = "2026-03-29T13:19:22.47Z" }, + { url = "https://files.pythonhosted.org/packages/c1/62/2b7a48fbb745d344742c0277f01286dead15f3f68e4f359fbfcf7b48f70f/numpy-2.4.4-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:e823b8b6edc81e747526f70f71a9c0a07ac4e7ad13020aa736bb7c9d67196115", size = 14694532, upload-time = "2026-03-29T13:19:25.581Z" }, + { url = "https://files.pythonhosted.org/packages/e5/87/499737bfba066b4a3bebff24a8f1c5b2dee410b209bc6668c9be692580f0/numpy-2.4.4-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:4a19d9dba1a76618dd86b164d608566f393f8ec6ac7c44f0cc879011c45e65af", size = 5199661, upload-time = "2026-03-29T13:19:28.31Z" }, + { url = "https://files.pythonhosted.org/packages/cd/da/464d551604320d1491bc345efed99b4b7034143a85787aab78d5691d5a0e/numpy-2.4.4-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:d2a8490669bfe99a233298348acc2d824d496dee0e66e31b66a6022c2ad74a5c", size = 6547539, upload-time = "2026-03-29T13:19:30.97Z" }, + { url = "https://files.pythonhosted.org/packages/7d/90/8d23e3b0dafd024bf31bdec225b3bb5c2dbfa6912f8a53b8659f21216cbf/numpy-2.4.4-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:45dbed2ab436a9e826e302fcdcbe9133f9b0006e5af7168afb8963a6520da103", size = 15668806, upload-time = "2026-03-29T13:19:33.887Z" }, + { url = "https://files.pythonhosted.org/packages/d1/73/a9d864e42a01896bb5974475438f16086be9ba1f0d19d0bb7a07427c4a8b/numpy-2.4.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c901b15172510173f5cb310eae652908340f8dede90fff9e3bf6c0d8dfd92f83", size = 16632682, upload-time = "2026-03-29T13:19:37.336Z" }, + { url = "https://files.pythonhosted.org/packages/34/fb/14570d65c3bde4e202a031210475ae9cde9b7686a2e7dc97ee67d2833b35/numpy-2.4.4-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:99d838547ace2c4aace6c4f76e879ddfe02bb58a80c1549928477862b7a6d6ed", size = 17019810, upload-time = "2026-03-29T13:19:40.963Z" }, + { url = "https://files.pythonhosted.org/packages/8a/77/2ba9d87081fd41f6d640c83f26fb7351e536b7ce6dd9061b6af5904e8e46/numpy-2.4.4-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:0aec54fd785890ecca25a6003fd9a5aed47ad607bbac5cd64f836ad8666f4959", size = 18357394, upload-time = "2026-03-29T13:19:44.859Z" }, + { url = "https://files.pythonhosted.org/packages/a2/23/52666c9a41708b0853fa3b1a12c90da38c507a3074883823126d4e9d5b30/numpy-2.4.4-cp313-cp313-win32.whl", hash = "sha256:07077278157d02f65c43b1b26a3886bce886f95d20aabd11f87932750dfb14ed", size = 5959556, upload-time = "2026-03-29T13:19:47.661Z" }, + { url = "https://files.pythonhosted.org/packages/57/fb/48649b4971cde70d817cf97a2a2fdc0b4d8308569f1dd2f2611959d2e0cf/numpy-2.4.4-cp313-cp313-win_amd64.whl", hash = "sha256:5c70f1cc1c4efbe316a572e2d8b9b9cc44e89b95f79ca3331553fbb63716e2bf", size = 12317311, upload-time = "2026-03-29T13:19:50.67Z" }, + { url = "https://files.pythonhosted.org/packages/ba/d8/11490cddd564eb4de97b4579ef6bfe6a736cc07e94c1598590ae25415e01/numpy-2.4.4-cp313-cp313-win_arm64.whl", hash = "sha256:ef4059d6e5152fa1a39f888e344c73fdc926e1b2dd58c771d67b0acfbf2aa67d", size = 10222060, upload-time = "2026-03-29T13:19:54.229Z" }, + { url = "https://files.pythonhosted.org/packages/99/5d/dab4339177a905aad3e2221c915b35202f1ec30d750dd2e5e9d9a72b804b/numpy-2.4.4-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:4bbc7f303d125971f60ec0aaad5e12c62d0d2c925f0ab1273debd0e4ba37aba5", size = 14822302, upload-time = "2026-03-29T13:19:57.585Z" }, + { url = "https://files.pythonhosted.org/packages/eb/e4/0564a65e7d3d97562ed6f9b0fd0fb0a6f559ee444092f105938b50043876/numpy-2.4.4-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:4d6d57903571f86180eb98f8f0c839fa9ebbfb031356d87f1361be91e433f5b7", size = 5327407, upload-time = "2026-03-29T13:20:00.601Z" }, + { url = "https://files.pythonhosted.org/packages/29/8d/35a3a6ce5ad371afa58b4700f1c820f8f279948cca32524e0a695b0ded83/numpy-2.4.4-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:4636de7fd195197b7535f231b5de9e4b36d2c440b6e566d2e4e4746e6af0ca93", size = 6647631, upload-time = "2026-03-29T13:20:02.855Z" }, + { url = "https://files.pythonhosted.org/packages/f4/da/477731acbd5a58a946c736edfdabb2ac5b34c3d08d1ba1a7b437fa0884df/numpy-2.4.4-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ad2e2ef14e0b04e544ea2fa0a36463f847f113d314aa02e5b402fdf910ef309e", size = 15727691, upload-time = "2026-03-29T13:20:06.004Z" }, + { url = "https://files.pythonhosted.org/packages/e6/db/338535d9b152beabeb511579598418ba0212ce77cf9718edd70262cc4370/numpy-2.4.4-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5a285b3b96f951841799528cd1f4f01cd70e7e0204b4abebac9463eecfcf2a40", size = 16681241, upload-time = "2026-03-29T13:20:09.417Z" }, + { url = "https://files.pythonhosted.org/packages/e2/a9/ad248e8f58beb7a0219b413c9c7d8151c5d285f7f946c3e26695bdbbe2df/numpy-2.4.4-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:f8474c4241bc18b750be2abea9d7a9ec84f46ef861dbacf86a4f6e043401f79e", size = 17085767, upload-time = "2026-03-29T13:20:13.126Z" }, + { url = "https://files.pythonhosted.org/packages/b5/1a/3b88ccd3694681356f70da841630e4725a7264d6a885c8d442a697e1146b/numpy-2.4.4-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:4e874c976154687c1f71715b034739b45c7711bec81db01914770373d125e392", size = 18403169, upload-time = "2026-03-29T13:20:17.096Z" }, + { url = "https://files.pythonhosted.org/packages/c2/c9/fcfd5d0639222c6eac7f304829b04892ef51c96a75d479214d77e3ce6e33/numpy-2.4.4-cp313-cp313t-win32.whl", hash = "sha256:9c585a1790d5436a5374bac930dad6ed244c046ed91b2b2a3634eb2971d21008", size = 6083477, upload-time = "2026-03-29T13:20:20.195Z" }, + { url = "https://files.pythonhosted.org/packages/d5/e3/3938a61d1c538aaec8ed6fd6323f57b0c2d2d2219512434c5c878db76553/numpy-2.4.4-cp313-cp313t-win_amd64.whl", hash = "sha256:93e15038125dc1e5345d9b5b68aa7f996ec33b98118d18c6ca0d0b7d6198b7e8", size = 12457487, upload-time = "2026-03-29T13:20:22.946Z" }, + { url = "https://files.pythonhosted.org/packages/97/6a/7e345032cc60501721ef94e0e30b60f6b0bd601f9174ebd36389a2b86d40/numpy-2.4.4-cp313-cp313t-win_arm64.whl", hash = "sha256:0dfd3f9d3adbe2920b68b5cd3d51444e13a10792ec7154cd0a2f6e74d4ab3233", size = 10292002, upload-time = "2026-03-29T13:20:25.909Z" }, +] + +[[package]] +name = "nvidia-cublas" +version = "13.1.0.3" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e1/a5/fce49e2ae977e0ccc084e5adafceb4f0ac0c8333cb6863501618a7277f67/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:c86fc7f7ae36d7528288c5d88098edcb7b02c633d262e7ddbb86b0ad91be5df2", size = 542851226, upload-time = "2025-10-09T08:59:04.818Z" }, + { url = "https://files.pythonhosted.org/packages/e7/44/423ac00af4dd95a5aeb27207e2c0d9b7118702149bf4704c3ddb55bb7429/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:ee8722c1f0145ab246bccb9e452153b5e0515fd094c3678df50b2a0888b8b171", size = 423133236, upload-time = "2025-10-09T08:59:32.536Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti" +version = "13.0.85" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/2a/80353b103fc20ce05ef51e928daed4b6015db4aaa9162ed0997090fe2250/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:796bd679890ee55fb14a94629b698b6db54bcfd833d391d5e94017dd9d7d3151", size = 10310827, upload-time = "2025-09-04T08:26:42.012Z" }, + { url = "https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8", size = 10715597, upload-time = "2025-09-04T08:26:51.312Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc" +version = "13.0.88" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575", size = 90215200, upload-time = "2025-09-04T08:28:44.204Z" }, + { url = "https://files.pythonhosted.org/packages/b7/dc/6bb80850e0b7edd6588d560758f17e0550893a1feaf436807d64d2da040f/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d27f20a0ca67a4bb34268a5e951033496c5b74870b868bacd046b1b8e0c3267b", size = 43015449, upload-time = "2025-09-04T08:28:20.239Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime" +version = "13.0.96" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/87/4f/17d7b9b8e285199c58ce28e31b5c5bbaa4d8271af06a89b6405258245de2/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ef9bcbe90493a2b9d810e43d249adb3d02e98dd30200d86607d8d02687c43f55", size = 2261060, upload-time = "2025-10-09T08:55:15.78Z" }, + { url = "https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548", size = 2243632, upload-time = "2025-10-09T08:55:36.117Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu13" +version = "9.19.0.56" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas", marker = "sys_platform != 'darwin'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/f1/84/26025437c1e6b61a707442184fa0c03d083b661adf3a3eecfd6d21677740/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:6ed29ffaee1176c612daf442e4dd6cfeb6a0caa43ddcbeb59da94953030b1be4", size = 433781201, upload-time = "2026-02-03T20:40:53.805Z" }, + { url = "https://files.pythonhosted.org/packages/a3/22/0b4b932655d17a6da1b92fa92ab12844b053bb2ac2475e179ba6f043da1e/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:d20e1734305e9d68889a96e3f35094d733ff1f83932ebe462753973e53a572bf", size = 366066321, upload-time = "2026-02-03T20:44:52.837Z" }, +] + +[[package]] +name = "nvidia-cufft" +version = "12.0.0.61" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink", marker = "sys_platform != 'darwin'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, + { url = "https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6c44f692dce8fd5ffd3e3df134b6cdb9c2f72d99cf40b62c32dde45eea9ddad3", size = 214085489, upload-time = "2025-09-04T08:31:56.044Z" }, +] + +[[package]] +name = "nvidia-cufile" +version = "1.15.1.6" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08a3ecefae5a01c7f5117351c64f17c7c62efa5fffdbe24fc7d298da19cd0b44", size = 1223672, upload-time = "2025-09-04T08:32:22.779Z" }, + { url = "https://files.pythonhosted.org/packages/ab/73/cc4a14c9813a8a0d509417cf5f4bdaba76e924d58beb9864f5a7baceefbf/nvidia_cufile-1.15.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:bdc0deedc61f548bddf7733bdc216456c2fdb101d020e1ab4b88d232d5e2f6d1", size = 1136992, upload-time = "2025-09-04T08:32:14.119Z" }, +] + +[[package]] +name = "nvidia-curand" +version = "10.4.0.35" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/72/7c2ae24fb6b63a32e6ae5d241cc65263ea18d08802aaae087d9f013335a2/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:133df5a7509c3e292aaa2b477afd0194f06ce4ea24d714d616ff36439cee349a", size = 61962106, upload-time = "2025-08-04T10:21:41.128Z" }, + { url = "https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:1aee33a5da6e1db083fe2b90082def8915f30f3248d5896bcec36a579d941bfc", size = 59544258, upload-time = "2025-08-04T10:22:03.992Z" }, +] + +[[package]] +name = "nvidia-cusolver" +version = "12.0.4.66" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas", marker = "sys_platform != 'darwin'" }, + { name = "nvidia-cusparse", marker = "sys_platform != 'darwin'" }, + { name = "nvidia-nvjitlink", marker = "sys_platform != 'darwin'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, + { url = "https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0a759da5dea5c0ea10fd307de75cdeb59e7ea4fcb8add0924859b944babf1112", size = 200941980, upload-time = "2025-09-04T08:33:22.767Z" }, +] + +[[package]] +name = "nvidia-cusparse" +version = "12.6.3.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink", marker = "sys_platform != 'darwin'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, + { url = "https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2b3c89c88d01ee0e477cb7f82ef60a11a4bcd57b6b87c33f789350b59759360b", size = 145942937, upload-time = "2025-09-04T08:33:58.029Z" }, +] + +[[package]] +name = "nvidia-cusparselt-cu13" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/10/8dcd1175260706a2fc92a16a52e306b71d4c1ea0b0cc4a9484183399818a/nvidia_cusparselt_cu13-0.8.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:400c6ed1cf6780fc6efedd64ec9f1345871767e6a1a0a552a1ea0578117ea77c", size = 220791277, upload-time = "2025-08-13T19:22:40.982Z" }, + { url = "https://files.pythonhosted.org/packages/fd/53/43b0d71f4e702fa9733f8b4571fdca50a8813f1e450b656c239beff12315/nvidia_cusparselt_cu13-0.8.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:25e30a8a7323935d4ad0340b95a0b69926eee755767e8e0b1cf8dd85b197d3fd", size = 169884119, upload-time = "2025-08-13T19:23:41.967Z" }, +] + +[[package]] +name = "nvidia-nccl-cu13" +version = "2.28.9" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/55/1920646a2e43ffd4fc958536b276197ed740e9e0c54105b4bb3521591fc7/nvidia_nccl_cu13-2.28.9-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:01c873ba1626b54caa12272ed228dc5b2781545e0ae8ba3f432a8ef1c6d78643", size = 196561677, upload-time = "2025-11-18T05:49:03.45Z" }, + { url = "https://files.pythonhosted.org/packages/b0/b4/878fefaad5b2bcc6fcf8d474a25e3e3774bc5133e4b58adff4d0bca238bc/nvidia_nccl_cu13-2.28.9-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:e4553a30f34195f3fa1da02a6da3d6337d28f2003943aa0a3d247bbc25fefc42", size = 196493177, upload-time = "2025-11-18T05:49:17.677Z" }, +] + +[[package]] +name = "nvidia-nvjitlink" +version = "13.0.88" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/7a/123e033aaff487c77107195fa5a2b8686795ca537935a24efae476c41f05/nvidia_nvjitlink-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:13a74f429e23b921c1109976abefacc69835f2f433ebd323d3946e11d804e47b", size = 40713933, upload-time = "2025-09-04T08:35:43.553Z" }, + { url = "https://files.pythonhosted.org/packages/ab/2c/93c5250e64df4f894f1cbb397c6fd71f79813f9fd79d7cd61de3f97b3c2d/nvidia_nvjitlink-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e931536ccc7d467a98ba1d8b89ff7fa7f1fa3b13f2b0069118cd7f47bff07d0c", size = 38768748, upload-time = "2025-09-04T08:35:20.008Z" }, +] + +[[package]] +name = "nvidia-nvshmem-cu13" +version = "3.4.5" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/0f/05cc9c720236dcd2db9c1ab97fff629e96821be2e63103569da0c9b72f19/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dc2a197f38e5d0376ad52cd1a2a3617d3cdc150fd5966f4aee9bcebb1d68fe9", size = 60215947, upload-time = "2025-09-06T00:32:20.022Z" }, + { url = "https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:290f0a2ee94c9f3687a02502f3b9299a9f9fe826e6d0287ee18482e78d495b80", size = 60412546, upload-time = "2025-09-06T00:32:41.564Z" }, +] + +[[package]] +name = "nvidia-nvtx" +version = "13.0.85" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4", size = 148047, upload-time = "2025-09-04T08:29:01.761Z" }, + { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" }, +] + +[[package]] +name = "onnx" +version = "1.21.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "protobuf" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c5/93/942d2a0f6a70538eea042ce0445c8aefd46559ad153469986f29a743c01c/onnx-1.21.0.tar.gz", hash = "sha256:4d8b67d0aaec5864c87633188b91cc520877477ec0254eda122bef8be43cd764", size = 12074608, upload-time = "2026-03-27T21:33:36.118Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7d/ae/cb644ec84c25e63575d9d8790fdcc5d1a11d67d3f62f872edb35fa38d158/onnx-1.21.0-cp312-abi3-macosx_12_0_universal2.whl", hash = "sha256:fc2635400fe39ff37ebc4e75342cc54450eadadf39c540ff132c319bf4960095", size = 17965930, upload-time = "2026-03-27T21:32:48.089Z" }, + { url = "https://files.pythonhosted.org/packages/6f/b6/eeb5903586645ef8a49b4b7892580438741acc3df91d7a5bd0f3a59ea9cb/onnx-1.21.0-cp312-abi3-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9003d5206c01fa2ff4b46311566865d8e493e1a6998d4009ec6de39843f1b59b", size = 17531344, upload-time = "2026-03-27T21:32:50.837Z" }, + { url = "https://files.pythonhosted.org/packages/a7/00/4823f06357892d1e60d6f34e7299d2ba4ed2108c487cc394f7ce85a3ff14/onnx-1.21.0-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a9261bd580fb8548c9c37b3c6750387eb8f21ea43c63880d37b2c622e1684285", size = 17613697, upload-time = "2026-03-27T21:32:54.222Z" }, + { url = "https://files.pythonhosted.org/packages/23/1d/391f3c567ae068c8ac4f1d1316bae97c9eb45e702f05975fe0e17ad441f0/onnx-1.21.0-cp312-abi3-win32.whl", hash = "sha256:9ea4e824964082811938a9250451d89c4ec474fe42dd36c038bfa5df31993d1e", size = 16287200, upload-time = "2026-03-27T21:32:57.277Z" }, + { url = "https://files.pythonhosted.org/packages/9c/a6/5eefbe5b40ea96de95a766bd2e0e751f35bdea2d4b951991ec9afaa69531/onnx-1.21.0-cp312-abi3-win_amd64.whl", hash = "sha256:458d91948ad9a7729a347550553b49ab6939f9af2cddf334e2116e45467dc61f", size = 16441045, upload-time = "2026-03-27T21:33:00.081Z" }, + { url = "https://files.pythonhosted.org/packages/63/c4/0ed8dc037a39113d2a4d66e0005e07751c299c46b993f1ad5c2c35664c20/onnx-1.21.0-cp312-abi3-win_arm64.whl", hash = "sha256:ca14bc4842fccc3187eb538f07eabeb25a779b39388b006db4356c07403a7bbb", size = 16403134, upload-time = "2026-03-27T21:33:03.987Z" }, + { url = "https://files.pythonhosted.org/packages/f8/89/0e1a9beb536401e2f45ac88735e123f2735e12fc7b56ff6c11727e097526/onnx-1.21.0-cp313-cp313t-macosx_12_0_universal2.whl", hash = "sha256:257d1d1deb6a652913698f1e3f33ef1ca0aa69174892fe38946d4572d89dd94f", size = 17975430, upload-time = "2026-03-27T21:33:07.005Z" }, + { url = "https://files.pythonhosted.org/packages/ec/46/e6dc71a7b3b317265591b20a5f71d0ff5c0d26c24e52283139dc90c66038/onnx-1.21.0-cp313-cp313t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7cd7cb8f6459311bdb557cbf6c0ccc6d8ace11c304d1bba0a30b4a4688e245f8", size = 17537435, upload-time = "2026-03-27T21:33:09.765Z" }, + { url = "https://files.pythonhosted.org/packages/49/2e/27affcac63eaf2ef183a44fd1a1354b11da64a6c72fe6f3fdcf5571bcee5/onnx-1.21.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7b58a4cfec8d9311b73dc083e4c1fa362069267881144c05139b3eba5dc3a840", size = 17617687, upload-time = "2026-03-27T21:33:12.619Z" }, + { url = "https://files.pythonhosted.org/packages/1c/5c/ac8ed15e941593a3672ce424280b764979026317811f2e8508432bfc3429/onnx-1.21.0-cp313-cp313t-win_amd64.whl", hash = "sha256:1a9baf882562c4cebf79589bebb7cd71a20e30b51158cac3e3bbaf27da6163bd", size = 16449402, upload-time = "2026-03-27T21:33:15.555Z" }, + { url = "https://files.pythonhosted.org/packages/0e/aa/d2231e0dcaad838217afc64c306c8152a080134d2034e247cc973d577674/onnx-1.21.0-cp313-cp313t-win_arm64.whl", hash = "sha256:bba12181566acf49b35875838eba49536a327b2944664b17125577d230c637ad", size = 16408273, upload-time = "2026-03-27T21:33:18.599Z" }, +] + +[[package]] +name = "onnx-ir" +version = "0.2.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "onnx" }, + { name = "sympy" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/35/e6/672fefb2f108d077f58181a7babf4c0f8d1182a30353ffc9c79c63afc5ee/onnx_ir-0.2.1.tar.gz", hash = "sha256:8b8b10a93f43e65962104de6070c43c5dacb0e3cdfefc7c8059dd83c9db64f35", size = 144279, upload-time = "2026-04-20T20:21:47.735Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8c/aa/f7a53321c60b9ad9ee184b6018292ed6b5389947592a2c8c09c736bb7f9e/onnx_ir-0.2.1-py3-none-any.whl", hash = "sha256:c7285da889312f91882de2092e298a9eeeefbfc1d1951c49d983992967eb09a7", size = 166792, upload-time = "2026-04-20T20:21:46.357Z" }, +] + +[[package]] +name = "onnxruntime" +version = "1.24.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "flatbuffers" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "protobuf" }, + { name = "sympy" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/d7/38/31db1b232b4ba960065a90c1506ad7a56995cd8482033184e97fadca17cc/onnxruntime-1.24.4-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:cad1c2b3f455c55678ab2a8caa51fb420c25e6e3cf10f4c23653cdabedc8de78", size = 17341875, upload-time = "2026-03-17T22:05:51.669Z" }, + { url = "https://files.pythonhosted.org/packages/aa/60/c4d1c8043eb42f8a9aa9e931c8c293d289c48ff463267130eca97d13357f/onnxruntime-1.24.4-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1a5c5a544b22f90859c88617ecb30e161ee3349fcc73878854f43d77f00558b5", size = 15172485, upload-time = "2026-03-17T22:03:32.182Z" }, + { url = "https://files.pythonhosted.org/packages/6d/ab/5b68110e0460d73fad814d5bd11c7b1ddcce5c37b10177eb264d6a36e331/onnxruntime-1.24.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0d640eb9f3782689b55cfa715094474cd5662f2f137be6a6f847a594b6e9705c", size = 17244912, upload-time = "2026-03-17T22:04:37.251Z" }, + { url = "https://files.pythonhosted.org/packages/8b/f4/6b89e297b93704345f0f3f8c62229bee323ef25682a3f9b4f89a39324950/onnxruntime-1.24.4-cp312-cp312-win_amd64.whl", hash = "sha256:535b29475ca42b593c45fbb2152fbf1cdf3f287315bf650e6a724a0a1d065cdb", size = 12596856, upload-time = "2026-03-17T22:05:41.224Z" }, + { url = "https://files.pythonhosted.org/packages/43/06/8b8ec6e9e6a474fcd5d772453f627ad4549dfe3ab8c0bf70af5afcde551b/onnxruntime-1.24.4-cp312-cp312-win_arm64.whl", hash = "sha256:e6214096e14b7b52e3bee1903dc12dc7ca09cb65e26664668a4620cc5e6f9a90", size = 12270275, upload-time = "2026-03-17T22:05:31.132Z" }, + { url = "https://files.pythonhosted.org/packages/e9/f0/8a21ec0a97e40abb7d8da1e8b20fb9e1af509cc6d191f6faa75f73622fb2/onnxruntime-1.24.4-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:e99a48078baaefa2b50fe5836c319499f71f13f76ed32d0211f39109147a49e0", size = 17341922, upload-time = "2026-03-17T22:03:56.364Z" }, + { url = "https://files.pythonhosted.org/packages/8b/25/d7908de8e08cee9abfa15b8aa82349b79733ae5865162a3609c11598805d/onnxruntime-1.24.4-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:dc4aaed1e5e1aaacf2343c838a30a7c3ade78f13eeb16817411f929d04040a13", size = 15172290, upload-time = "2026-03-17T22:03:37.124Z" }, + { url = "https://files.pythonhosted.org/packages/7f/72/105ec27a78c5aa0154a7c0cd8c41c19a97799c3b12fc30392928997e3be3/onnxruntime-1.24.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e30c972bc02e072911aabb6891453ec73795386c0af2b761b65444b8a4c4745f", size = 17244738, upload-time = "2026-03-17T22:04:40.625Z" }, + { url = "https://files.pythonhosted.org/packages/05/fb/a592736d968c2f58e12de4d52088dda8e0e724b26ad5c0487263adb45875/onnxruntime-1.24.4-cp313-cp313-win_amd64.whl", hash = "sha256:3b6ba8b0181a3aa88edab00eb01424ffc06f42e71095a91186c2249415fcff93", size = 12597435, upload-time = "2026-03-17T22:05:43.826Z" }, + { url = "https://files.pythonhosted.org/packages/ad/04/ae2479e9841b64bd2eb44f8a64756c62593f896514369a11243b1b86ca5c/onnxruntime-1.24.4-cp313-cp313-win_arm64.whl", hash = "sha256:71d6a5c1821d6e8586a024000ece458db8f2fc0ecd050435d45794827ce81e19", size = 12269852, upload-time = "2026-03-17T22:05:33.353Z" }, + { url = "https://files.pythonhosted.org/packages/b4/af/a479a536c4398ffaf49fbbe755f45d5b8726bdb4335ab31b537f3d7149b8/onnxruntime-1.24.4-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1700f559c8086d06b2a4d5de51e62cb4ff5e2631822f71a36db8c72383db71ee", size = 15176861, upload-time = "2026-03-17T22:03:40.143Z" }, + { url = "https://files.pythonhosted.org/packages/be/13/19f5da70c346a76037da2c2851ecbf1266e61d7f0dcdb887c667210d4608/onnxruntime-1.24.4-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4c74e268dc808e61e63784d43f9ddcdaf50a776c2819e8bd1d1b11ef64bf7e36", size = 17247454, upload-time = "2026-03-17T22:04:46.643Z" }, +] + +[[package]] +name = "onnxscript" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "onnx" }, + { name = "onnx-ir" }, + { name = "packaging" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9b/99/fd948eba63ba65b52265a4cd09a14f96bb9f5b730fcef58876c4358bf406/onnxscript-0.7.0.tar.gz", hash = "sha256:c95ed7b339b02cface56ee27689565c46612e1fc542c562298dddfdad5268dc5", size = 612032, upload-time = "2026-04-20T17:09:19.775Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b9/ce/2ed92575cc3be4ea1db5f38f16f20765f9b20b69b14d6c1d9972658a8ee9/onnxscript-0.7.0-py3-none-any.whl", hash = "sha256:5b356907d4501e9919f8599c91d8da967406a37b1fac2b40caa55a49acf242ea", size = 714842, upload-time = "2026-04-20T17:09:22.089Z" }, +] + +[[package]] +name = "orjson" +version = "3.11.8" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9d/1b/2024d06792d0779f9dbc51531b61c24f76c75b9f4ce05e6f3377a1814cea/orjson-3.11.8.tar.gz", hash = "sha256:96163d9cdc5a202703e9ad1b9ae757d5f0ca62f4fa0cc93d1f27b0e180cc404e", size = 5603832, upload-time = "2026-03-31T16:16:27.878Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/01/f6/8d58b32ab32d9215973a1688aebd098252ee8af1766c0e4e36e7831f0295/orjson-3.11.8-cp312-cp312-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:1cd0b77e77c95758f8e1100139844e99f3ccc87e71e6fc8e1c027e55807c549f", size = 229233, upload-time = "2026-03-31T16:15:12.762Z" }, + { url = "https://files.pythonhosted.org/packages/a9/8b/2ffe35e71f6b92622e8ea4607bf33ecf7dfb51b3619dcfabfd36cbe2d0a5/orjson-3.11.8-cp312-cp312-macosx_15_0_arm64.whl", hash = "sha256:6a3d159d5ffa0e3961f353c4b036540996bf8b9697ccc38261c0eac1fd3347a6", size = 128772, upload-time = "2026-03-31T16:15:14.237Z" }, + { url = "https://files.pythonhosted.org/packages/27/d2/1f8682ae50d5c6897a563cb96bc106da8c9cb5b7b6e81a52e4cc086679b9/orjson-3.11.8-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:76070a76e9c5ae661e2d9848f216980d8d533e0f8143e6ed462807b242e3c5e8", size = 131946, upload-time = "2026-03-31T16:15:15.607Z" }, + { url = "https://files.pythonhosted.org/packages/52/4b/5500f76f0eece84226e0689cb48dcde081104c2fa6e2483d17ca13685ffb/orjson-3.11.8-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:54153d21520a71a4c82a0dbb4523e468941d549d221dc173de0f019678cf3813", size = 130368, upload-time = "2026-03-31T16:15:17.066Z" }, + { url = "https://files.pythonhosted.org/packages/da/4e/58b927e08fbe9840e6c920d9e299b051ea667463b1f39a56e668669f8508/orjson-3.11.8-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:469ac2125611b7c5741a0b3798cd9e5786cbad6345f9f400c77212be89563bec", size = 135540, upload-time = "2026-03-31T16:15:18.404Z" }, + { url = "https://files.pythonhosted.org/packages/56/7c/ba7cb871cba1bcd5cd02ee34f98d894c6cea96353ad87466e5aef2429c60/orjson-3.11.8-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:14778ffd0f6896aa613951a7fbf4690229aa7a543cb2bfbe9f358e08aafa9546", size = 146877, upload-time = "2026-03-31T16:15:19.833Z" }, + { url = "https://files.pythonhosted.org/packages/0b/5d/eb9c25fc1386696c6a342cd361c306452c75e0b55e86ad602dd4827a7fd7/orjson-3.11.8-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:ea56a955056a6d6c550cf18b3348656a9d9a4f02e2d0c02cabf3c73f1055d506", size = 132837, upload-time = "2026-03-31T16:15:21.282Z" }, + { url = "https://files.pythonhosted.org/packages/37/87/5ddeb7fc1fbd9004aeccab08426f34c81a5b4c25c7061281862b015fce2b/orjson-3.11.8-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:53a0f57e59a530d18a142f4d4ba6dfc708dc5fdedce45e98ff06b44930a2a48f", size = 133624, upload-time = "2026-03-31T16:15:22.641Z" }, + { url = "https://files.pythonhosted.org/packages/22/09/90048793db94ee4b2fcec4ac8e5ddb077367637d6650be896b3494b79bb7/orjson-3.11.8-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:9b48e274f8824567d74e2158199e269597edf00823a1b12b63d48462bbf5123e", size = 141904, upload-time = "2026-03-31T16:15:24.435Z" }, + { url = "https://files.pythonhosted.org/packages/c0/cf/eb284847487821a5d415e54149a6449ba9bfc5872ce63ab7be41b8ec401c/orjson-3.11.8-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:3f262401086a3960586af06c054609365e98407151f5ea24a62893a40d80dbbb", size = 423742, upload-time = "2026-03-31T16:15:26.155Z" }, + { url = "https://files.pythonhosted.org/packages/44/09/e12423d327071c851c13e76936f144a96adacfc037394dec35ac3fc8d1e8/orjson-3.11.8-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:8e8c6218b614badf8e229b697865df4301afa74b791b6c9ade01d19a9953a942", size = 147806, upload-time = "2026-03-31T16:15:27.909Z" }, + { url = "https://files.pythonhosted.org/packages/b3/6d/37c2589ba864e582ffe7611643314785c6afb1f83c701654ef05daa8fcc7/orjson-3.11.8-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:093d489fa039ddade2db541097dbb484999fcc65fc2b0ff9819141e2ab364f25", size = 136485, upload-time = "2026-03-31T16:15:29.749Z" }, + { url = "https://files.pythonhosted.org/packages/be/c9/135194a02ab76b04ed9a10f68624b7ebd238bbe55548878b11ff15a0f352/orjson-3.11.8-cp312-cp312-win32.whl", hash = "sha256:e0950ed1bcb9893f4293fd5c5a7ee10934fbf82c4101c70be360db23ce24b7d2", size = 131966, upload-time = "2026-03-31T16:15:31.687Z" }, + { url = "https://files.pythonhosted.org/packages/ed/9a/9796f8fbe3cf30ce9cb696748dbb535e5c87be4bf4fe2e9ca498ef1fa8cf/orjson-3.11.8-cp312-cp312-win_amd64.whl", hash = "sha256:3cf17c141617b88ced4536b2135c552490f07799f6ad565948ea07bef0dcb9a6", size = 127441, upload-time = "2026-03-31T16:15:33.333Z" }, + { url = "https://files.pythonhosted.org/packages/cc/47/5aaf54524a7a4a0dd09dd778f3fa65dd2108290615b652e23d944152bc8e/orjson-3.11.8-cp312-cp312-win_arm64.whl", hash = "sha256:48854463b0572cc87dac7d981aa72ed8bf6deedc0511853dc76b8bbd5482d36d", size = 127364, upload-time = "2026-03-31T16:15:34.748Z" }, + { url = "https://files.pythonhosted.org/packages/66/7f/95fba509bb2305fab0073558f1e8c3a2ec4b2afe58ed9fcb7d3b8beafe94/orjson-3.11.8-cp313-cp313-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:3f23426851d98478c8970da5991f84784a76682213cd50eb73a1da56b95239dc", size = 229180, upload-time = "2026-03-31T16:15:36.426Z" }, + { url = "https://files.pythonhosted.org/packages/f6/9d/b237215c743ca073697d759b5503abd2cb8a0d7b9c9e21f524bcf176ab66/orjson-3.11.8-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:ebaed4cef74a045b83e23537b52ef19a367c7e3f536751e355a2a394f8648559", size = 128754, upload-time = "2026-03-31T16:15:38.049Z" }, + { url = "https://files.pythonhosted.org/packages/42/3d/27d65b6d11e63f133781425f132807aef793ed25075fec686fc8e46dd528/orjson-3.11.8-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:97c8f5d3b62380b70c36ffacb2a356b7c6becec86099b177f73851ba095ef623", size = 131877, upload-time = "2026-03-31T16:15:39.484Z" }, + { url = "https://files.pythonhosted.org/packages/dd/cc/faee30cd8f00421999e40ef0eba7332e3a625ce91a58200a2f52c7fef235/orjson-3.11.8-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:436c4922968a619fb7fef1ccd4b8b3a76c13b67d607073914d675026e911a65c", size = 130361, upload-time = "2026-03-31T16:15:41.274Z" }, + { url = "https://files.pythonhosted.org/packages/5c/bb/a6c55896197f97b6d4b4e7c7fd77e7235517c34f5d6ad5aadd43c54c6d7c/orjson-3.11.8-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1ab359aff0436d80bfe8a23b46b5fea69f1e18aaf1760a709b4787f1318b317f", size = 135521, upload-time = "2026-03-31T16:15:42.758Z" }, + { url = "https://files.pythonhosted.org/packages/9c/7c/ca3a3525aa32ff636ebb1778e77e3587b016ab2edb1b618b36ba96f8f2c0/orjson-3.11.8-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f89b6d0b3a8d81e1929d3ab3d92bbc225688bd80a770c49432543928fe09ac55", size = 146862, upload-time = "2026-03-31T16:15:44.341Z" }, + { url = "https://files.pythonhosted.org/packages/3c/0c/18a9d7f18b5edd37344d1fd5be17e94dc652c67826ab749c6e5948a78112/orjson-3.11.8-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:29c009e7a2ca9ad0ed1376ce20dd692146a5d9fe4310848904b6b4fee5c5c137", size = 132847, upload-time = "2026-03-31T16:15:46.368Z" }, + { url = "https://files.pythonhosted.org/packages/23/91/7e722f352ad67ca573cee44de2a58fb810d0f4eb4e33276c6a557979fd8a/orjson-3.11.8-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:705b895b781b3e395c067129d8551655642dfe9437273211d5404e87ac752b53", size = 133637, upload-time = "2026-03-31T16:15:48.123Z" }, + { url = "https://files.pythonhosted.org/packages/af/04/32845ce13ac5bd1046ddb02ac9432ba856cc35f6d74dde95864fe0ad5523/orjson-3.11.8-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:88006eda83858a9fdf73985ce3804e885c2befb2f506c9a3723cdeb5a2880e3e", size = 141906, upload-time = "2026-03-31T16:15:49.626Z" }, + { url = "https://files.pythonhosted.org/packages/02/5e/c551387ddf2d7106d9039369862245c85738b828844d13b99ccb8d61fd06/orjson-3.11.8-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:55120759e61309af7fcf9e961c6f6af3dde5921cdb3ee863ef63fd9db126cae6", size = 423722, upload-time = "2026-03-31T16:15:51.176Z" }, + { url = "https://files.pythonhosted.org/packages/00/a3/ecfe62434096f8a794d4976728cb59bcfc4a643977f21c2040545d37eb4c/orjson-3.11.8-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:98bdc6cb889d19bed01de46e67574a2eab61f5cc6b768ed50e8ac68e9d6ffab6", size = 147801, upload-time = "2026-03-31T16:15:52.939Z" }, + { url = "https://files.pythonhosted.org/packages/18/6d/0dce10b9f6643fdc59d99333871a38fa5a769d8e2fc34a18e5d2bfdee900/orjson-3.11.8-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:708c95f925a43ab9f34625e45dcdadf09ec8a6e7b664a938f2f8d5650f6c090b", size = 136460, upload-time = "2026-03-31T16:15:54.431Z" }, + { url = "https://files.pythonhosted.org/packages/01/d6/6dde4f31842d87099238f1f07b459d24edc1a774d20687187443ab044191/orjson-3.11.8-cp313-cp313-win32.whl", hash = "sha256:01c4e5a6695dc09098f2e6468a251bc4671c50922d4d745aff1a0a33a0cf5b8d", size = 131956, upload-time = "2026-03-31T16:15:56.081Z" }, + { url = "https://files.pythonhosted.org/packages/c1/f9/4e494a56e013db957fb77186b818b916d4695b8fa2aa612364974160e91b/orjson-3.11.8-cp313-cp313-win_amd64.whl", hash = "sha256:c154a35dd1330707450bb4d4e7dd1f17fa6f42267a40c1e8a1daa5e13719b4b8", size = 127410, upload-time = "2026-03-31T16:15:57.54Z" }, + { url = "https://files.pythonhosted.org/packages/57/7f/803203d00d6edb6e9e7eef421d4e1adbb5ea973e40b3533f3cfd9aeb374e/orjson-3.11.8-cp313-cp313-win_arm64.whl", hash = "sha256:4861bde57f4d253ab041e374f44023460e60e71efaa121f3c5f0ed457c3a701e", size = 127338, upload-time = "2026-03-31T16:15:59.106Z" }, +] + +[[package]] +name = "packaging" +version = "26.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/df/de/0d2b39fb4af88a0258f3bac87dfcbb48e73fbdea4a2ed0e2213f9a4c2f9a/packaging-26.1.tar.gz", hash = "sha256:f042152b681c4bfac5cae2742a55e103d27ab2ec0f3d88037136b6bfe7c9c5de", size = 215519, upload-time = "2026-04-14T21:12:49.362Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7a/c2/920ef838e2f0028c8262f16101ec09ebd5969864e5a64c4c05fad0617c56/packaging-26.1-py3-none-any.whl", hash = "sha256:5d9c0669c6285e491e0ced2eee587eaf67b670d94a19e94e3984a481aba6802f", size = 95831, upload-time = "2026-04-14T21:12:47.56Z" }, +] + +[[package]] +name = "paramiko" +version = "4.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "bcrypt" }, + { name = "cryptography" }, + { name = "invoke" }, + { name = "pynacl" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1f/e7/81fdcbc7f190cdb058cffc9431587eb289833bdd633e2002455ca9bb13d4/paramiko-4.0.0.tar.gz", hash = "sha256:6a25f07b380cc9c9a88d2b920ad37167ac4667f8d9886ccebd8f90f654b5d69f", size = 1630743, upload-time = "2025-08-04T01:02:03.711Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a9/90/a744336f5af32c433bd09af7854599682a383b37cfd78f7de263de6ad6cb/paramiko-4.0.0-py3-none-any.whl", hash = "sha256:0e20e00ac666503bf0b4eda3b6d833465a2b7aff2e2b3d79a8bba5ef144ee3b9", size = 223932, upload-time = "2025-08-04T01:02:02.029Z" }, +] + +[[package]] +name = "pillow" +version = "11.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f3/0d/d0d6dea55cd152ce3d6767bb38a8fc10e33796ba4ba210cbab9354b6d238/pillow-11.3.0.tar.gz", hash = "sha256:3828ee7586cd0b2091b6209e5ad53e20d0649bbe87164a459d0676e035e8f523", size = 47113069, upload-time = "2025-07-01T09:16:30.666Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/40/fe/1bc9b3ee13f68487a99ac9529968035cca2f0a51ec36892060edcc51d06a/pillow-11.3.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:fdae223722da47b024b867c1ea0be64e0df702c5e0a60e27daad39bf960dd1e4", size = 5278800, upload-time = "2025-07-01T09:14:17.648Z" }, + { url = "https://files.pythonhosted.org/packages/2c/32/7e2ac19b5713657384cec55f89065fb306b06af008cfd87e572035b27119/pillow-11.3.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:921bd305b10e82b4d1f5e802b6850677f965d8394203d182f078873851dada69", size = 4686296, upload-time = "2025-07-01T09:14:19.828Z" }, + { url = "https://files.pythonhosted.org/packages/8e/1e/b9e12bbe6e4c2220effebc09ea0923a07a6da1e1f1bfbc8d7d29a01ce32b/pillow-11.3.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:eb76541cba2f958032d79d143b98a3a6b3ea87f0959bbe256c0b5e416599fd5d", size = 5871726, upload-time = "2025-07-03T13:10:04.448Z" }, + { url = "https://files.pythonhosted.org/packages/8d/33/e9200d2bd7ba00dc3ddb78df1198a6e80d7669cce6c2bdbeb2530a74ec58/pillow-11.3.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:67172f2944ebba3d4a7b54f2e95c786a3a50c21b88456329314caaa28cda70f6", size = 7644652, upload-time = "2025-07-03T13:10:10.391Z" }, + { url = "https://files.pythonhosted.org/packages/41/f1/6f2427a26fc683e00d985bc391bdd76d8dd4e92fac33d841127eb8fb2313/pillow-11.3.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:97f07ed9f56a3b9b5f49d3661dc9607484e85c67e27f3e8be2c7d28ca032fec7", size = 5977787, upload-time = "2025-07-01T09:14:21.63Z" }, + { url = "https://files.pythonhosted.org/packages/e4/c9/06dd4a38974e24f932ff5f98ea3c546ce3f8c995d3f0985f8e5ba48bba19/pillow-11.3.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:676b2815362456b5b3216b4fd5bd89d362100dc6f4945154ff172e206a22c024", size = 6645236, upload-time = "2025-07-01T09:14:23.321Z" }, + { url = "https://files.pythonhosted.org/packages/40/e7/848f69fb79843b3d91241bad658e9c14f39a32f71a301bcd1d139416d1be/pillow-11.3.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3e184b2f26ff146363dd07bde8b711833d7b0202e27d13540bfe2e35a323a809", size = 6086950, upload-time = "2025-07-01T09:14:25.237Z" }, + { url = "https://files.pythonhosted.org/packages/0b/1a/7cff92e695a2a29ac1958c2a0fe4c0b2393b60aac13b04a4fe2735cad52d/pillow-11.3.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:6be31e3fc9a621e071bc17bb7de63b85cbe0bfae91bb0363c893cbe67247780d", size = 6723358, upload-time = "2025-07-01T09:14:27.053Z" }, + { url = "https://files.pythonhosted.org/packages/26/7d/73699ad77895f69edff76b0f332acc3d497f22f5d75e5360f78cbcaff248/pillow-11.3.0-cp312-cp312-win32.whl", hash = "sha256:7b161756381f0918e05e7cb8a371fff367e807770f8fe92ecb20d905d0e1c149", size = 6275079, upload-time = "2025-07-01T09:14:30.104Z" }, + { url = "https://files.pythonhosted.org/packages/8c/ce/e7dfc873bdd9828f3b6e5c2bbb74e47a98ec23cc5c74fc4e54462f0d9204/pillow-11.3.0-cp312-cp312-win_amd64.whl", hash = "sha256:a6444696fce635783440b7f7a9fc24b3ad10a9ea3f0ab66c5905be1c19ccf17d", size = 6986324, upload-time = "2025-07-01T09:14:31.899Z" }, + { url = "https://files.pythonhosted.org/packages/16/8f/b13447d1bf0b1f7467ce7d86f6e6edf66c0ad7cf44cf5c87a37f9bed9936/pillow-11.3.0-cp312-cp312-win_arm64.whl", hash = "sha256:2aceea54f957dd4448264f9bf40875da0415c83eb85f55069d89c0ed436e3542", size = 2423067, upload-time = "2025-07-01T09:14:33.709Z" }, + { url = "https://files.pythonhosted.org/packages/1e/93/0952f2ed8db3a5a4c7a11f91965d6184ebc8cd7cbb7941a260d5f018cd2d/pillow-11.3.0-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:1c627742b539bba4309df89171356fcb3cc5a9178355b2727d1b74a6cf155fbd", size = 2128328, upload-time = "2025-07-01T09:14:35.276Z" }, + { url = "https://files.pythonhosted.org/packages/4b/e8/100c3d114b1a0bf4042f27e0f87d2f25e857e838034e98ca98fe7b8c0a9c/pillow-11.3.0-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:30b7c02f3899d10f13d7a48163c8969e4e653f8b43416d23d13d1bbfdc93b9f8", size = 2170652, upload-time = "2025-07-01T09:14:37.203Z" }, + { url = "https://files.pythonhosted.org/packages/aa/86/3f758a28a6e381758545f7cdb4942e1cb79abd271bea932998fc0db93cb6/pillow-11.3.0-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:7859a4cc7c9295f5838015d8cc0a9c215b77e43d07a25e460f35cf516df8626f", size = 2227443, upload-time = "2025-07-01T09:14:39.344Z" }, + { url = "https://files.pythonhosted.org/packages/01/f4/91d5b3ffa718df2f53b0dc109877993e511f4fd055d7e9508682e8aba092/pillow-11.3.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:ec1ee50470b0d050984394423d96325b744d55c701a439d2bd66089bff963d3c", size = 5278474, upload-time = "2025-07-01T09:14:41.843Z" }, + { url = "https://files.pythonhosted.org/packages/f9/0e/37d7d3eca6c879fbd9dba21268427dffda1ab00d4eb05b32923d4fbe3b12/pillow-11.3.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:7db51d222548ccfd274e4572fdbf3e810a5e66b00608862f947b163e613b67dd", size = 4686038, upload-time = "2025-07-01T09:14:44.008Z" }, + { url = "https://files.pythonhosted.org/packages/ff/b0/3426e5c7f6565e752d81221af9d3676fdbb4f352317ceafd42899aaf5d8a/pillow-11.3.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2d6fcc902a24ac74495df63faad1884282239265c6839a0a6416d33faedfae7e", size = 5864407, upload-time = "2025-07-03T13:10:15.628Z" }, + { url = "https://files.pythonhosted.org/packages/fc/c1/c6c423134229f2a221ee53f838d4be9d82bab86f7e2f8e75e47b6bf6cd77/pillow-11.3.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f0f5d8f4a08090c6d6d578351a2b91acf519a54986c055af27e7a93feae6d3f1", size = 7639094, upload-time = "2025-07-03T13:10:21.857Z" }, + { url = "https://files.pythonhosted.org/packages/ba/c9/09e6746630fe6372c67c648ff9deae52a2bc20897d51fa293571977ceb5d/pillow-11.3.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c37d8ba9411d6003bba9e518db0db0c58a680ab9fe5179f040b0463644bc9805", size = 5973503, upload-time = "2025-07-01T09:14:45.698Z" }, + { url = "https://files.pythonhosted.org/packages/d5/1c/a2a29649c0b1983d3ef57ee87a66487fdeb45132df66ab30dd37f7dbe162/pillow-11.3.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:13f87d581e71d9189ab21fe0efb5a23e9f28552d5be6979e84001d3b8505abe8", size = 6642574, upload-time = "2025-07-01T09:14:47.415Z" }, + { url = "https://files.pythonhosted.org/packages/36/de/d5cc31cc4b055b6c6fd990e3e7f0f8aaf36229a2698501bcb0cdf67c7146/pillow-11.3.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:023f6d2d11784a465f09fd09a34b150ea4672e85fb3d05931d89f373ab14abb2", size = 6084060, upload-time = "2025-07-01T09:14:49.636Z" }, + { url = "https://files.pythonhosted.org/packages/d5/ea/502d938cbaeec836ac28a9b730193716f0114c41325db428e6b280513f09/pillow-11.3.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:45dfc51ac5975b938e9809451c51734124e73b04d0f0ac621649821a63852e7b", size = 6721407, upload-time = "2025-07-01T09:14:51.962Z" }, + { url = "https://files.pythonhosted.org/packages/45/9c/9c5e2a73f125f6cbc59cc7087c8f2d649a7ae453f83bd0362ff7c9e2aee2/pillow-11.3.0-cp313-cp313-win32.whl", hash = "sha256:a4d336baed65d50d37b88ca5b60c0fa9d81e3a87d4a7930d3880d1624d5b31f3", size = 6273841, upload-time = "2025-07-01T09:14:54.142Z" }, + { url = "https://files.pythonhosted.org/packages/23/85/397c73524e0cd212067e0c969aa245b01d50183439550d24d9f55781b776/pillow-11.3.0-cp313-cp313-win_amd64.whl", hash = "sha256:0bce5c4fd0921f99d2e858dc4d4d64193407e1b99478bc5cacecba2311abde51", size = 6978450, upload-time = "2025-07-01T09:14:56.436Z" }, + { url = "https://files.pythonhosted.org/packages/17/d2/622f4547f69cd173955194b78e4d19ca4935a1b0f03a302d655c9f6aae65/pillow-11.3.0-cp313-cp313-win_arm64.whl", hash = "sha256:1904e1264881f682f02b7f8167935cce37bc97db457f8e7849dc3a6a52b99580", size = 2423055, upload-time = "2025-07-01T09:14:58.072Z" }, + { url = "https://files.pythonhosted.org/packages/dd/80/a8a2ac21dda2e82480852978416cfacd439a4b490a501a288ecf4fe2532d/pillow-11.3.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:4c834a3921375c48ee6b9624061076bc0a32a60b5532b322cc0ea64e639dd50e", size = 5281110, upload-time = "2025-07-01T09:14:59.79Z" }, + { url = "https://files.pythonhosted.org/packages/44/d6/b79754ca790f315918732e18f82a8146d33bcd7f4494380457ea89eb883d/pillow-11.3.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:5e05688ccef30ea69b9317a9ead994b93975104a677a36a8ed8106be9260aa6d", size = 4689547, upload-time = "2025-07-01T09:15:01.648Z" }, + { url = "https://files.pythonhosted.org/packages/49/20/716b8717d331150cb00f7fdd78169c01e8e0c219732a78b0e59b6bdb2fd6/pillow-11.3.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1019b04af07fc0163e2810167918cb5add8d74674b6267616021ab558dc98ced", size = 5901554, upload-time = "2025-07-03T13:10:27.018Z" }, + { url = "https://files.pythonhosted.org/packages/74/cf/a9f3a2514a65bb071075063a96f0a5cf949c2f2fce683c15ccc83b1c1cab/pillow-11.3.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f944255db153ebb2b19c51fe85dd99ef0ce494123f21b9db4877ffdfc5590c7c", size = 7669132, upload-time = "2025-07-03T13:10:33.01Z" }, + { url = "https://files.pythonhosted.org/packages/98/3c/da78805cbdbee9cb43efe8261dd7cc0b4b93f2ac79b676c03159e9db2187/pillow-11.3.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1f85acb69adf2aaee8b7da124efebbdb959a104db34d3a2cb0f3793dbae422a8", size = 6005001, upload-time = "2025-07-01T09:15:03.365Z" }, + { url = "https://files.pythonhosted.org/packages/6c/fa/ce044b91faecf30e635321351bba32bab5a7e034c60187fe9698191aef4f/pillow-11.3.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:05f6ecbeff5005399bb48d198f098a9b4b6bdf27b8487c7f38ca16eeb070cd59", size = 6668814, upload-time = "2025-07-01T09:15:05.655Z" }, + { url = "https://files.pythonhosted.org/packages/7b/51/90f9291406d09bf93686434f9183aba27b831c10c87746ff49f127ee80cb/pillow-11.3.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:a7bc6e6fd0395bc052f16b1a8670859964dbd7003bd0af2ff08342eb6e442cfe", size = 6113124, upload-time = "2025-07-01T09:15:07.358Z" }, + { url = "https://files.pythonhosted.org/packages/cd/5a/6fec59b1dfb619234f7636d4157d11fb4e196caeee220232a8d2ec48488d/pillow-11.3.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:83e1b0161c9d148125083a35c1c5a89db5b7054834fd4387499e06552035236c", size = 6747186, upload-time = "2025-07-01T09:15:09.317Z" }, + { url = "https://files.pythonhosted.org/packages/49/6b/00187a044f98255225f172de653941e61da37104a9ea60e4f6887717e2b5/pillow-11.3.0-cp313-cp313t-win32.whl", hash = "sha256:2a3117c06b8fb646639dce83694f2f9eac405472713fcb1ae887469c0d4f6788", size = 6277546, upload-time = "2025-07-01T09:15:11.311Z" }, + { url = "https://files.pythonhosted.org/packages/e8/5c/6caaba7e261c0d75bab23be79f1d06b5ad2a2ae49f028ccec801b0e853d6/pillow-11.3.0-cp313-cp313t-win_amd64.whl", hash = "sha256:857844335c95bea93fb39e0fa2726b4d9d758850b34075a7e3ff4f4fa3aa3b31", size = 6985102, upload-time = "2025-07-01T09:15:13.164Z" }, + { url = "https://files.pythonhosted.org/packages/f3/7e/b623008460c09a0cb38263c93b828c666493caee2eb34ff67f778b87e58c/pillow-11.3.0-cp313-cp313t-win_arm64.whl", hash = "sha256:8797edc41f3e8536ae4b10897ee2f637235c94f27404cac7297f7b607dd0716e", size = 2424803, upload-time = "2025-07-01T09:15:15.695Z" }, +] + +[[package]] +name = "platformdirs" +version = "4.9.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9f/4a/0883b8e3802965322523f0b200ecf33d31f10991d0401162f4b23c698b42/platformdirs-4.9.6.tar.gz", hash = "sha256:3bfa75b0ad0db84096ae777218481852c0ebc6c727b3168c1b9e0118e458cf0a", size = 29400, upload-time = "2026-04-09T00:04:10.812Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/75/a6/a0a304dc33b49145b21f4808d763822111e67d1c3a32b524a1baf947b6e1/platformdirs-4.9.6-py3-none-any.whl", hash = "sha256:e61adb1d5e5cb3441b4b7710bea7e4c12250ca49439228cc1021c00dcfac0917", size = 21348, upload-time = "2026-04-09T00:04:09.463Z" }, +] + +[[package]] +name = "prettytable" +version = "3.17.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wcwidth" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/79/45/b0847d88d6cfeb4413566738c8bbf1e1995fad3d42515327ff32cc1eb578/prettytable-3.17.0.tar.gz", hash = "sha256:59f2590776527f3c9e8cf9fe7b66dd215837cca96a9c39567414cbc632e8ddb0", size = 67892, upload-time = "2025-11-14T17:33:20.212Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ee/8c/83087ebc47ab0396ce092363001fa37c17153119ee282700c0713a195853/prettytable-3.17.0-py3-none-any.whl", hash = "sha256:aad69b294ddbe3e1f95ef8886a060ed1666a0b83018bbf56295f6f226c43d287", size = 34433, upload-time = "2025-11-14T17:33:19.093Z" }, +] + +[[package]] +name = "proglog" +version = "0.1.12" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "tqdm" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c2/af/c108866c452eda1132f3d6b3cb6be2ae8430c97e9309f38ca9dbd430af37/proglog-0.1.12.tar.gz", hash = "sha256:361ee074721c277b89b75c061336cb8c5f287c92b043efa562ccf7866cda931c", size = 8794, upload-time = "2025-05-09T14:36:18.316Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/1b/f7ea6cde25621cd9236541c66ff018f4268012a534ec31032bcb187dc5e7/proglog-0.1.12-py3-none-any.whl", hash = "sha256:ccaafce51e80a81c65dc907a460c07ccb8ec1f78dc660cfd8f9ec3a22f01b84c", size = 6337, upload-time = "2025-05-09T14:36:16.798Z" }, +] + +[[package]] +name = "protobuf" +version = "7.34.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6b/6b/a0e95cad1ad7cc3f2c6821fcab91671bd5b78bd42afb357bb4765f29bc41/protobuf-7.34.1.tar.gz", hash = "sha256:9ce42245e704cc5027be797c1db1eb93184d44d1cdd71811fb2d9b25ad541280", size = 454708, upload-time = "2026-03-20T17:34:47.036Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/11/3325d41e6ee15bf1125654301211247b042563bcc898784351252549a8ad/protobuf-7.34.1-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:d8b2cc79c4d8f62b293ad9b11ec3aebce9af481fa73e64556969f7345ebf9fc7", size = 429247, upload-time = "2026-03-20T17:34:37.024Z" }, + { url = "https://files.pythonhosted.org/packages/eb/9d/aa69df2724ff63efa6f72307b483ce0827f4347cc6d6df24b59e26659fef/protobuf-7.34.1-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:5185e0e948d07abe94bb76ec9b8416b604cfe5da6f871d67aad30cbf24c3110b", size = 325753, upload-time = "2026-03-20T17:34:38.751Z" }, + { url = "https://files.pythonhosted.org/packages/92/e8/d174c91fd48e50101943f042b09af9029064810b734e4160bbe282fa1caa/protobuf-7.34.1-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:403b093a6e28a960372b44e5eb081775c9b056e816a8029c61231743d63f881a", size = 340198, upload-time = "2026-03-20T17:34:39.871Z" }, + { url = "https://files.pythonhosted.org/packages/53/1b/3b431694a4dc6d37b9f653f0c64b0a0d9ec074ee810710c0c3da21d67ba7/protobuf-7.34.1-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:8ff40ce8cd688f7265326b38d5a1bed9bfdf5e6723d49961432f83e21d5713e4", size = 324267, upload-time = "2026-03-20T17:34:41.1Z" }, + { url = "https://files.pythonhosted.org/packages/85/29/64de04a0ac142fb685fd09999bc3d337943fb386f3a0ec57f92fd8203f97/protobuf-7.34.1-cp310-abi3-win32.whl", hash = "sha256:34b84ce27680df7cca9f231043ada0daa55d0c44a2ddfaa58ec1d0d89d8bf60a", size = 426628, upload-time = "2026-03-20T17:34:42.536Z" }, + { url = "https://files.pythonhosted.org/packages/4d/87/cb5e585192a22b8bd457df5a2c16a75ea0db9674c3a0a39fc9347d84e075/protobuf-7.34.1-cp310-abi3-win_amd64.whl", hash = "sha256:e97b55646e6ce5cbb0954a8c28cd39a5869b59090dfaa7df4598a7fba869468c", size = 437901, upload-time = "2026-03-20T17:34:44.112Z" }, + { url = "https://files.pythonhosted.org/packages/88/95/608f665226bca68b736b79e457fded9a2a38c4f4379a4a7614303d9db3bc/protobuf-7.34.1-py3-none-any.whl", hash = "sha256:bb3812cd53aefea2b028ef42bd780f5b96407247f20c6ef7c679807e9d188f11", size = 170715, upload-time = "2026-03-20T17:34:45.384Z" }, +] + +[[package]] +name = "pyarrow" +version = "24.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/91/13/13e1069b351bdc3881266e11147ffccf687505dbb0ea74036237f5d454a5/pyarrow-24.0.0.tar.gz", hash = "sha256:85fe721a14dd823aca09127acbb06c3ca723efbd436c004f16bca601b04dcc83", size = 1180261, upload-time = "2026-04-21T10:51:25.837Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b4/a9/9686d9f07837f91f775e8932659192e02c74f9d8920524b480b85212cc68/pyarrow-24.0.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:6233c9ed9ab9d1db47de57d9753256d9dcffbf42db341576099f0fd9f6bf4810", size = 34981559, upload-time = "2026-04-21T10:47:22.17Z" }, + { url = "https://files.pythonhosted.org/packages/80/b6/0ddf0e9b6ead3474ab087ae598c76b031fc45532bf6a63f3a553440fb258/pyarrow-24.0.0-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:f7616236ec1bc2b15bfdec22a71ab38851c86f8f05ff64f379e1278cf20c634a", size = 36663654, upload-time = "2026-04-21T10:47:28.315Z" }, + { url = "https://files.pythonhosted.org/packages/7c/3b/926382efe8ce27ba729071d3566ade6dfb86bdf112f366000196b2f5780a/pyarrow-24.0.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:1617043b99bd33e5318ae18eb2919af09c71322ef1ca46566cdafc6e6712fb66", size = 45679394, upload-time = "2026-04-21T10:47:34.821Z" }, + { url = "https://files.pythonhosted.org/packages/b3/7a/829f7d9dfd37c207206081d6dad474d81dde29952401f07f2ba507814818/pyarrow-24.0.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:6165461f55ef6314f026de6638d661188e3455d3ec49834556a0ebbdbace18bb", size = 48863122, upload-time = "2026-04-21T10:47:42.056Z" }, + { url = "https://files.pythonhosted.org/packages/5f/e8/f88ce625fe8babaae64e8db2d417c7653adb3019b08aae85c5ed787dc816/pyarrow-24.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3b13dedfe76a0ad2d1d859b0811b53827a4e9d93a0bcb05cf59333ab4980cc7e", size = 49376032, upload-time = "2026-04-21T10:47:48.967Z" }, + { url = "https://files.pythonhosted.org/packages/36/7a/82c363caa145fff88fb475da50d3bf52bb024f61917be5424c3392eaf878/pyarrow-24.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:25ea65d868eb04015cd18e6df2fbe98f07e5bda2abefabcb88fce39a947716f6", size = 51929490, upload-time = "2026-04-21T10:47:55.981Z" }, + { url = "https://files.pythonhosted.org/packages/66/1c/e3e72c8014ad2743ca64a701652c733cc5cbcee15c0463a32a8c55518d9e/pyarrow-24.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:295f0a7f2e242dabd513737cf076007dc5b2d59237e3eca37b05c0c6446f3826", size = 27355660, upload-time = "2026-04-21T10:48:01.718Z" }, + { url = "https://files.pythonhosted.org/packages/6f/d3/a1abf004482026ddc17f4503db227787fa3cfe41ec5091ff20e4fea55e57/pyarrow-24.0.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:02b001b3ed4723caa44f6cd1af2d5c86aa2cf9971dacc2ffa55b21237713dfba", size = 34976759, upload-time = "2026-04-21T10:48:07.258Z" }, + { url = "https://files.pythonhosted.org/packages/4f/4a/34f0a36d28a2dd32225301b79daad44e243dc1a2bb77d43b60749be255c4/pyarrow-24.0.0-cp313-cp313-macosx_12_0_x86_64.whl", hash = "sha256:04920d6a71aabd08a0417709efce97d45ea8e6fb733d9ca9ecffb13c67839f68", size = 36658471, upload-time = "2026-04-21T10:48:13.347Z" }, + { url = "https://files.pythonhosted.org/packages/1f/78/543b94712ae8bb1a6023bcc1acf1a740fbff8286747c289cd9468fced2a5/pyarrow-24.0.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:a964266397740257f16f7bb2e4f08a0c81454004beab8ff59dd531b73610e9f2", size = 45675981, upload-time = "2026-04-21T10:48:20.201Z" }, + { url = "https://files.pythonhosted.org/packages/84/9f/8fb7c222b100d314137fa40ec050de56cd8c6d957d1cfff685ce72f15b17/pyarrow-24.0.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6f066b179d68c413374294bc1735f68475457c933258df594443bb9d88ddc2a0", size = 48859172, upload-time = "2026-04-21T10:48:27.541Z" }, + { url = "https://files.pythonhosted.org/packages/a7/d3/1ea72538e6c8b3b475ed78d1049a2c518e655761ea50fe1171fc855fcab7/pyarrow-24.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:1183baeb14c5f587b1ec52831e665718ce632caab84b7cd6b85fd44f96114495", size = 49385733, upload-time = "2026-04-21T10:48:34.7Z" }, + { url = "https://files.pythonhosted.org/packages/c3/be/c3d8b06a1ba35f2260f8e1f771abbee7d5e345c0937aab90675706b1690a/pyarrow-24.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:806f24b4085453c197a5078218d1ee08783ebbba271badd153d1ae22a3ee804f", size = 51934335, upload-time = "2026-04-21T10:48:42.099Z" }, + { url = "https://files.pythonhosted.org/packages/9c/62/89e07a1e7329d2cde3e3c6994ba0839a24977a2beda8be6005ea3d860b99/pyarrow-24.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:e4505fc6583f7b05ab854934896bcac8253b04ac1171a77dfb73efef92076d91", size = 27271748, upload-time = "2026-04-21T10:49:42.532Z" }, + { url = "https://files.pythonhosted.org/packages/17/1a/cff3a59f80b5b1658549d46611b67163f65e0664431c076ad728bf9d5af4/pyarrow-24.0.0-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:1a4e45017efbf115032e4475ee876d525e0e36c742214fbe405332480ecd6275", size = 35238554, upload-time = "2026-04-21T10:48:48.526Z" }, + { url = "https://files.pythonhosted.org/packages/a8/99/cce0f42a327bfef2c420fb6078a3eb834826e5d6697bf3009fe11d2ad051/pyarrow-24.0.0-cp313-cp313t-macosx_12_0_x86_64.whl", hash = "sha256:7986f1fa71cee060ad00758bcc79d3a93bab8559bf978fab9e53472a2e25a17b", size = 36782301, upload-time = "2026-04-21T10:48:55.181Z" }, + { url = "https://files.pythonhosted.org/packages/2a/66/8e560d5ff6793ca29aca213c53eec0dd482dd46cb93b2819e5aab52e4252/pyarrow-24.0.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:d3e0b61e8efb24ed38898e5cdc5fffa9124be480008d401a1f8071500494ae42", size = 45721929, upload-time = "2026-04-21T10:49:03.676Z" }, + { url = "https://files.pythonhosted.org/packages/27/0c/a26e25505d030716e078d9f16eb74973cbf0b33b672884e9f9da1c83b871/pyarrow-24.0.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:55a3bc1e3df3b5567b7d27ef551b2283f0c68a5e86f1cd56abc569da4f31335b", size = 48825365, upload-time = "2026-04-21T10:49:11.714Z" }, + { url = "https://files.pythonhosted.org/packages/5f/eb/771f9ecb0c65e73fe9dccdd1717901b9594f08c4515d000c7c62df573811/pyarrow-24.0.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:641f795b361874ac9da5294f8f443dfdbee355cf2bd9e3b8d97aaac2306b9b37", size = 49451819, upload-time = "2026-04-21T10:49:21.474Z" }, + { url = "https://files.pythonhosted.org/packages/48/da/61ae89a88732f5a785646f3ec6125dbb640fa98a540eb2b9889caa561403/pyarrow-24.0.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8adc8e6ce5fccf5dc707046ae4914fd537def529709cc0d285d37a7f9cd442ca", size = 51909252, upload-time = "2026-04-21T10:49:31.164Z" }, + { url = "https://files.pythonhosted.org/packages/cb/1a/8dd5cafab7b66573fa91c03d06d213356ad4edd71813aa75e08ce2b3a844/pyarrow-24.0.0-cp313-cp313t-win_amd64.whl", hash = "sha256:9b18371ad2f44044b81a8d23bc2d8a9b6a6226dca775e8e16cfee640473d6c5d", size = 27388127, upload-time = "2026-04-21T10:49:37.334Z" }, +] + +[[package]] +name = "pybullet" +version = "3.2.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/de/95/b9a98cd4ed948d4f25a7d0b13cb0c1f87e58dc6d113be9e18940641eb25f/pybullet-3.2.7.tar.gz", hash = "sha256:042879db8d101ac7590dee475fc6aded508b85fe1273fdbbfde1d88bd200e14f", size = 80508379, upload-time = "2025-01-30T00:34:00.527Z" } + +[[package]] +name = "pycollada" +version = "0.9.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "python-dateutil" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5a/8d/52a5364a17eb96129962cae8d3ee7658775e085ad0ba38388684ad5944e9/pycollada-0.9.3.tar.gz", hash = "sha256:c34d6dcf0fe2eba5896f71c96d37a1c0fe1a61f08440fa0cfcec3dc2895d3302", size = 110826, upload-time = "2026-01-24T15:45:23.625Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/07/86/f1f61b7a0701f9d1299e5293d083318019f91021a4d449f94d59dbe024e4/pycollada-0.9.3-py3-none-any.whl", hash = "sha256:636e6496f60987586db82455ea7bbd9ade775e8181c6590c83b698b6cd53a9f5", size = 129206, upload-time = "2026-01-24T15:45:22.182Z" }, +] + +[[package]] +name = "pycparser" +version = "3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1b/7d/92392ff7815c21062bea51aa7b87d45576f649f16458d78b7cf94b9ab2e6/pycparser-3.0.tar.gz", hash = "sha256:600f49d217304a5902ac3c37e1281c9fe94e4d0489de643a9504c5cdfdfc6b29", size = 103492, upload-time = "2026-01-21T14:26:51.89Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0c/c3/44f3fbbfa403ea2a7c779186dc20772604442dde72947e7d01069cbe98e3/pycparser-3.0-py3-none-any.whl", hash = "sha256:b727414169a36b7d524c1c3e31839a521725078d7b2ff038656844266160a992", size = 48172, upload-time = "2026-01-21T14:26:50.693Z" }, +] + +[[package]] +name = "pydantic" +version = "2.13.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-types" }, + { name = "pydantic-core" }, + { name = "typing-extensions" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d9/e4/40d09941a2cebcb20609b86a559817d5b9291c49dd6f8c87e5feffbe703a/pydantic-2.13.3.tar.gz", hash = "sha256:af09e9d1d09f4e7fe37145c1f577e1d61ceb9a41924bf0094a36506285d0a84d", size = 844068, upload-time = "2026-04-20T14:46:43.632Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f3/0a/fd7d723f8f8153418fb40cf9c940e82004fce7e987026b08a68a36dd3fe7/pydantic-2.13.3-py3-none-any.whl", hash = "sha256:6db14ac8dfc9a1e57f87ea2c0de670c251240f43cb0c30a5130e9720dc612927", size = 471981, upload-time = "2026-04-20T14:46:41.402Z" }, +] + +[[package]] +name = "pydantic-core" +version = "2.46.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2a/ef/f7abb56c49382a246fd2ce9c799691e3c3e7175ec74b14d99e798bcddb1a/pydantic_core-2.46.3.tar.gz", hash = "sha256:41c178f65b8c29807239d47e6050262eb6bf84eb695e41101e62e38df4a5bc2c", size = 471412, upload-time = "2026-04-20T14:40:56.672Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4b/cb/5b47425556ecc1f3fe18ed2a0083188aa46e1dd812b06e406475b3a5d536/pydantic_core-2.46.3-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:b11b59b3eee90a80a36701ddb4576d9ae31f93f05cb9e277ceaa09e6bf074a67", size = 2101946, upload-time = "2026-04-20T14:40:52.581Z" }, + { url = "https://files.pythonhosted.org/packages/a1/4f/2fb62c2267cae99b815bbf4a7b9283812c88ca3153ef29f7707200f1d4e5/pydantic_core-2.46.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:af8653713055ea18a3abc1537fe2ebc42f5b0bbb768d1eb79fd74eb47c0ac089", size = 1951612, upload-time = "2026-04-20T14:42:42.996Z" }, + { url = "https://files.pythonhosted.org/packages/50/6e/b7348fd30d6556d132cddd5bd79f37f96f2601fe0608afac4f5fb01ec0b3/pydantic_core-2.46.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:75a519dab6d63c514f3a81053e5266c549679e4aa88f6ec57f2b7b854aceb1b0", size = 1977027, upload-time = "2026-04-20T14:42:02.001Z" }, + { url = "https://files.pythonhosted.org/packages/82/11/31d60ee2b45540d3fb0b29302a393dbc01cd771c473f5b5147bcd353e593/pydantic_core-2.46.3-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:a6cd87cb1575b1ad05ba98894c5b5c96411ef678fa2f6ed2576607095b8d9789", size = 2063008, upload-time = "2026-04-20T14:44:17.952Z" }, + { url = "https://files.pythonhosted.org/packages/8a/db/3a9d1957181b59258f44a2300ab0f0be9d1e12d662a4f57bb31250455c52/pydantic_core-2.46.3-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f80a55484b8d843c8ada81ebf70a682f3f00a3d40e378c06cf17ecb44d280d7d", size = 2233082, upload-time = "2026-04-20T14:40:57.934Z" }, + { url = "https://files.pythonhosted.org/packages/9c/e1/3277c38792aeb5cfb18c2f0c5785a221d9ff4e149abbe1184d53d5f72273/pydantic_core-2.46.3-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:3861f1731b90c50a3266316b9044f5c9b405eecb8e299b0a7120596334e4fe9c", size = 2304615, upload-time = "2026-04-20T14:42:12.584Z" }, + { url = "https://files.pythonhosted.org/packages/5e/d5/e3d9717c9eba10855325650afd2a9cba8e607321697f18953af9d562da2f/pydantic_core-2.46.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fb528e295ed31570ac3dcc9bfdd6e0150bc11ce6168ac87a8082055cf1a67395", size = 2094380, upload-time = "2026-04-20T14:43:05.522Z" }, + { url = "https://files.pythonhosted.org/packages/a1/20/abac35dedcbfd66c6f0b03e4e3564511771d6c9b7ede10a362d03e110d9b/pydantic_core-2.46.3-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:367508faa4973b992b271ba1494acaab36eb7e8739d1e47be5035fb1ea225396", size = 2135429, upload-time = "2026-04-20T14:41:55.549Z" }, + { url = "https://files.pythonhosted.org/packages/6c/a5/41bfd1df69afad71b5cf0535055bccc73022715ad362edbc124bc1e021d7/pydantic_core-2.46.3-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:5ad3c826fe523e4becf4fe39baa44286cff85ef137c729a2c5e269afbfd0905d", size = 2174582, upload-time = "2026-04-20T14:41:45.96Z" }, + { url = "https://files.pythonhosted.org/packages/79/65/38d86ea056b29b2b10734eb23329b7a7672ca604df4f2b6e9c02d4ee22fe/pydantic_core-2.46.3-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:ec638c5d194ef8af27db69f16c954a09797c0dc25015ad6123eb2c73a4d271ca", size = 2187533, upload-time = "2026-04-20T14:40:55.367Z" }, + { url = "https://files.pythonhosted.org/packages/b6/55/a1129141678a2026badc539ad1dee0a71d06f54c2f06a4bd68c030ac781b/pydantic_core-2.46.3-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:28ed528c45446062ee66edb1d33df5d88828ae167de76e773a3c7f64bd14e976", size = 2332985, upload-time = "2026-04-20T14:44:13.05Z" }, + { url = "https://files.pythonhosted.org/packages/d7/60/cb26f4077719f709e54819f4e8e1d43f4091f94e285eb6bd21e1190a7b7c/pydantic_core-2.46.3-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:aed19d0c783886d5bd86d80ae5030006b45e28464218747dcf83dabfdd092c7b", size = 2373670, upload-time = "2026-04-20T14:41:53.421Z" }, + { url = "https://files.pythonhosted.org/packages/6b/7e/c3f21882bdf1d8d086876f81b5e296206c69c6082551d776895de7801fa0/pydantic_core-2.46.3-cp312-cp312-win32.whl", hash = "sha256:06d5d8820cbbdb4147578c1fe7ffcd5b83f34508cb9f9ab76e807be7db6ff0a4", size = 1966722, upload-time = "2026-04-20T14:44:30.588Z" }, + { url = "https://files.pythonhosted.org/packages/57/be/6b5e757b859013ebfbd7adba02f23b428f37c86dcbf78b5bb0b4ffd36e99/pydantic_core-2.46.3-cp312-cp312-win_amd64.whl", hash = "sha256:c3212fda0ee959c1dd04c60b601ec31097aaa893573a3a1abd0a47bcac2968c1", size = 2072970, upload-time = "2026-04-20T14:42:54.248Z" }, + { url = "https://files.pythonhosted.org/packages/bf/f8/a989b21cc75e9a32d24192ef700eea606521221a89faa40c919ce884f2b1/pydantic_core-2.46.3-cp312-cp312-win_arm64.whl", hash = "sha256:f1f8338dd7a7f31761f1f1a3c47503a9a3b34eea3c8b01fa6ee96408affb5e72", size = 2035963, upload-time = "2026-04-20T14:44:20.4Z" }, + { url = "https://files.pythonhosted.org/packages/9b/3c/9b5e8eb9821936d065439c3b0fb1490ffa64163bfe7e1595985a47896073/pydantic_core-2.46.3-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:12bc98de041458b80c86c56b24df1d23832f3e166cbaff011f25d187f5c62c37", size = 2102109, upload-time = "2026-04-20T14:41:24.219Z" }, + { url = "https://files.pythonhosted.org/packages/91/97/1c41d1f5a19f241d8069f1e249853bcce378cdb76eec8ab636d7bc426280/pydantic_core-2.46.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:85348b8f89d2c3508b65b16c3c33a4da22b8215138d8b996912bb1532868885f", size = 1951820, upload-time = "2026-04-20T14:42:14.236Z" }, + { url = "https://files.pythonhosted.org/packages/30/b4/d03a7ae14571bc2b6b3c7b122441154720619afe9a336fa3a95434df5e2f/pydantic_core-2.46.3-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1105677a6df914b1fb71a81b96c8cce7726857e1717d86001f29be06a25ee6f8", size = 1977785, upload-time = "2026-04-20T14:42:31.648Z" }, + { url = "https://files.pythonhosted.org/packages/ae/0c/4086f808834b59e3c8f1aa26df8f4b6d998cdcf354a143d18ef41529d1fe/pydantic_core-2.46.3-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:87082cd65669a33adeba5470769e9704c7cf026cc30afb9cc77fd865578ebaad", size = 2062761, upload-time = "2026-04-20T14:40:37.093Z" }, + { url = "https://files.pythonhosted.org/packages/fa/71/a649be5a5064c2df0db06e0a512c2281134ed2fcc981f52a657936a7527c/pydantic_core-2.46.3-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:60e5f66e12c4f5212d08522963380eaaeac5ebd795826cfd19b2dfb0c7a52b9c", size = 2232989, upload-time = "2026-04-20T14:42:59.254Z" }, + { url = "https://files.pythonhosted.org/packages/a2/84/7756e75763e810b3a710f4724441d1ecc5883b94aacb07ca71c5fb5cfb69/pydantic_core-2.46.3-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b6cdf19bf84128d5e7c37e8a73a0c5c10d51103a650ac585d42dd6ae233f2b7f", size = 2303975, upload-time = "2026-04-20T14:41:32.287Z" }, + { url = "https://files.pythonhosted.org/packages/6c/35/68a762e0c1e31f35fa0dac733cbd9f5b118042853698de9509c8e5bf128b/pydantic_core-2.46.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:031bb17f4885a43773c8c763089499f242aee2ea85cf17154168775dccdecf35", size = 2095325, upload-time = "2026-04-20T14:42:47.685Z" }, + { url = "https://files.pythonhosted.org/packages/77/bf/1bf8c9a8e91836c926eae5e3e51dce009bf495a60ca56060689d3df3f340/pydantic_core-2.46.3-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:bcf2a8b2982a6673693eae7348ef3d8cf3979c1d63b54fca7c397a635cc68687", size = 2133368, upload-time = "2026-04-20T14:41:22.766Z" }, + { url = "https://files.pythonhosted.org/packages/e5/50/87d818d6bab915984995157ceb2380f5aac4e563dddbed6b56f0ed057aba/pydantic_core-2.46.3-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:28e8cf2f52d72ced402a137145923a762cbb5081e48b34312f7a0c8f55928ec3", size = 2173908, upload-time = "2026-04-20T14:42:52.044Z" }, + { url = "https://files.pythonhosted.org/packages/91/88/a311fb306d0bd6185db41fa14ae888fb81d0baf648a761ae760d30819d33/pydantic_core-2.46.3-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:17eaface65d9fc5abb940003020309c1bf7a211f5f608d7870297c367e6f9022", size = 2186422, upload-time = "2026-04-20T14:43:29.55Z" }, + { url = "https://files.pythonhosted.org/packages/8f/79/28fd0d81508525ab2054fef7c77a638c8b5b0afcbbaeee493cf7c3fef7e1/pydantic_core-2.46.3-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:93fd339f23408a07e98950a89644f92c54d8729719a40b30c0a30bb9ebc55d23", size = 2332709, upload-time = "2026-04-20T14:42:16.134Z" }, + { url = "https://files.pythonhosted.org/packages/b3/21/795bf5fe5c0f379308b8ef19c50dedab2e7711dbc8d0c2acf08f1c7daa05/pydantic_core-2.46.3-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:23cbdb3aaa74dfe0837975dbf69b469753bbde8eacace524519ffdb6b6e89eb7", size = 2372428, upload-time = "2026-04-20T14:41:10.974Z" }, + { url = "https://files.pythonhosted.org/packages/45/b3/ed14c659cbe7605e3ef063077680a64680aec81eb1a04763a05190d49b7f/pydantic_core-2.46.3-cp313-cp313-win32.whl", hash = "sha256:610eda2e3838f401105e6326ca304f5da1e15393ae25dacae5c5c63f2c275b13", size = 1965601, upload-time = "2026-04-20T14:41:42.128Z" }, + { url = "https://files.pythonhosted.org/packages/ef/bb/adb70d9a762ddd002d723fbf1bd492244d37da41e3af7b74ad212609027e/pydantic_core-2.46.3-cp313-cp313-win_amd64.whl", hash = "sha256:68cc7866ed863db34351294187f9b729964c371ba33e31c26f478471c52e1ed0", size = 2071517, upload-time = "2026-04-20T14:43:36.096Z" }, + { url = "https://files.pythonhosted.org/packages/52/eb/66faefabebfe68bd7788339c9c9127231e680b11906368c67ce112fdb47f/pydantic_core-2.46.3-cp313-cp313-win_arm64.whl", hash = "sha256:f64b5537ac62b231572879cd08ec05600308636a5d63bcbdb15063a466977bec", size = 2035802, upload-time = "2026-04-20T14:43:38.507Z" }, + { url = "https://files.pythonhosted.org/packages/34/42/f426db557e8ab2791bc7562052299944a118655496fbff99914e564c0a94/pydantic_core-2.46.3-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:b12dd51f1187c2eb489af8e20f880362db98e954b54ab792fa5d92e8bcc6b803", size = 2091877, upload-time = "2026-04-20T14:43:27.091Z" }, + { url = "https://files.pythonhosted.org/packages/5c/4f/86a832a9d14df58e663bfdf4627dc00d3317c2bd583c4fb23390b0f04b8e/pydantic_core-2.46.3-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:f00a0961b125f1a47af7bcc17f00782e12f4cd056f83416006b30111d941dfa3", size = 1932428, upload-time = "2026-04-20T14:40:45.781Z" }, + { url = "https://files.pythonhosted.org/packages/11/1a/fe857968954d93fb78e0d4b6df5c988c74c4aaa67181c60be7cfe327c0ca/pydantic_core-2.46.3-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:57697d7c056aca4bbb680200f96563e841a6386ac1129370a0102592f4dddff5", size = 1997550, upload-time = "2026-04-20T14:44:02.425Z" }, + { url = "https://files.pythonhosted.org/packages/17/eb/9d89ad2d9b0ba8cd65393d434471621b98912abb10fbe1df08e480ba57b5/pydantic_core-2.46.3-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd35aa21299def8db7ef4fe5c4ff862941a9a158ca7b63d61e66fe67d30416b4", size = 2137657, upload-time = "2026-04-20T14:42:45.149Z" }, +] + +[[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 = "pynacl" +version = "1.6.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cffi", marker = "platform_python_implementation != 'PyPy'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d9/9a/4019b524b03a13438637b11538c82781a5eda427394380381af8f04f467a/pynacl-1.6.2.tar.gz", hash = "sha256:018494d6d696ae03c7e656e5e74cdfd8ea1326962cc401bcf018f1ed8436811c", size = 3511692, upload-time = "2026-01-01T17:48:10.851Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/be/7b/4845bbf88e94586ec47a432da4e9107e3fc3ce37eb412b1398630a37f7dd/pynacl-1.6.2-cp38-abi3-macosx_10_10_universal2.whl", hash = "sha256:c949ea47e4206af7c8f604b8278093b674f7c79ed0d4719cc836902bf4517465", size = 388458, upload-time = "2026-01-01T17:32:16.829Z" }, + { url = "https://files.pythonhosted.org/packages/1e/b4/e927e0653ba63b02a4ca5b4d852a8d1d678afbf69b3dbf9c4d0785ac905c/pynacl-1.6.2-cp38-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8845c0631c0be43abdd865511c41eab235e0be69c81dc66a50911594198679b0", size = 800020, upload-time = "2026-01-01T17:32:18.34Z" }, + { url = "https://files.pythonhosted.org/packages/7f/81/d60984052df5c97b1d24365bc1e30024379b42c4edcd79d2436b1b9806f2/pynacl-1.6.2-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:22de65bb9010a725b0dac248f353bb072969c94fa8d6b1f34b87d7953cf7bbe4", size = 1399174, upload-time = "2026-01-01T17:32:20.239Z" }, + { url = "https://files.pythonhosted.org/packages/68/f7/322f2f9915c4ef27d140101dd0ed26b479f7e6f5f183590fd32dfc48c4d3/pynacl-1.6.2-cp38-abi3-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:46065496ab748469cdd999246d17e301b2c24ae2fdf739132e580a0e94c94a87", size = 835085, upload-time = "2026-01-01T17:32:22.24Z" }, + { url = "https://files.pythonhosted.org/packages/3e/d0/f301f83ac8dbe53442c5a43f6a39016f94f754d7a9815a875b65e218a307/pynacl-1.6.2-cp38-abi3-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8a66d6fb6ae7661c58995f9c6435bda2b1e68b54b598a6a10247bfcdadac996c", size = 1437614, upload-time = "2026-01-01T17:32:23.766Z" }, + { url = "https://files.pythonhosted.org/packages/c4/58/fc6e649762b029315325ace1a8c6be66125e42f67416d3dbd47b69563d61/pynacl-1.6.2-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:26bfcd00dcf2cf160f122186af731ae30ab120c18e8375684ec2670dccd28130", size = 818251, upload-time = "2026-01-01T17:32:25.69Z" }, + { url = "https://files.pythonhosted.org/packages/c9/a8/b917096b1accc9acd878819a49d3d84875731a41eb665f6ebc826b1af99e/pynacl-1.6.2-cp38-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:c8a231e36ec2cab018c4ad4358c386e36eede0319a0c41fed24f840b1dac59f6", size = 1402859, upload-time = "2026-01-01T17:32:27.215Z" }, + { url = "https://files.pythonhosted.org/packages/85/42/fe60b5f4473e12c72f977548e4028156f4d340b884c635ec6b063fe7e9a5/pynacl-1.6.2-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:68be3a09455743ff9505491220b64440ced8973fe930f270c8e07ccfa25b1f9e", size = 791926, upload-time = "2026-01-01T17:32:29.314Z" }, + { url = "https://files.pythonhosted.org/packages/fa/f9/e40e318c604259301cc091a2a63f237d9e7b424c4851cafaea4ea7c4834e/pynacl-1.6.2-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8b097553b380236d51ed11356c953bf8ce36a29a3e596e934ecabe76c985a577", size = 1363101, upload-time = "2026-01-01T17:32:31.263Z" }, + { url = "https://files.pythonhosted.org/packages/48/47/e761c254f410c023a469284a9bc210933e18588ca87706ae93002c05114c/pynacl-1.6.2-cp38-abi3-win32.whl", hash = "sha256:5811c72b473b2f38f7e2a3dc4f8642e3a3e9b5e7317266e4ced1fba85cae41aa", size = 227421, upload-time = "2026-01-01T17:32:33.076Z" }, + { url = "https://files.pythonhosted.org/packages/41/ad/334600e8cacc7d86587fe5f565480fde569dfb487389c8e1be56ac21d8ac/pynacl-1.6.2-cp38-abi3-win_amd64.whl", hash = "sha256:62985f233210dee6548c223301b6c25440852e13d59a8b81490203c3227c5ba0", size = 239754, upload-time = "2026-01-01T17:32:34.557Z" }, + { url = "https://files.pythonhosted.org/packages/29/7d/5945b5af29534641820d3bd7b00962abbbdfee84ec7e19f0d5b3175f9a31/pynacl-1.6.2-cp38-abi3-win_arm64.whl", hash = "sha256:834a43af110f743a754448463e8fd61259cd4ab5bbedcf70f9dabad1d28a394c", size = 184801, upload-time = "2026-01-01T17:32:36.309Z" }, +] + +[[package]] +name = "pynput" +version = "1.8.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "evdev", marker = "'linux' in sys_platform" }, + { name = "pyobjc-framework-applicationservices", marker = "sys_platform == 'darwin'" }, + { name = "pyobjc-framework-quartz", marker = "sys_platform == 'darwin'" }, + { name = "python-xlib", marker = "'linux' in sys_platform" }, + { name = "six" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f0/c3/dccf44c68225046df5324db0cc7d563a560635355b3e5f1d249468268a6f/pynput-1.8.1.tar.gz", hash = "sha256:70d7c8373ee98911004a7c938742242840a5628c004573d84ba849d4601df81e", size = 82289, upload-time = "2025-03-17T17:12:01.481Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/59/4f/ac3fa906ae8a375a536b12794128c5efacade9eaa917a35dfd27ce0c7400/pynput-1.8.1-py2.py3-none-any.whl", hash = "sha256:42dfcf27404459ca16ca889c8fb8ffe42a9fe54f722fd1a3e130728e59e768d2", size = 91693, upload-time = "2025-03-17T17:12:00.094Z" }, +] + +[[package]] +name = "pyobjc-core" +version = "12.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b8/b6/d5612eb40be4fd5ef88c259339e6313f46ba67577a95d86c3470b951fce0/pyobjc_core-12.1.tar.gz", hash = "sha256:2bb3903f5387f72422145e1466b3ac3f7f0ef2e9960afa9bcd8961c5cbf8bd21", size = 1000532, upload-time = "2025-11-14T10:08:28.292Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/64/5a/6b15e499de73050f4a2c88fff664ae154307d25dc04da8fb38998a428358/pyobjc_core-12.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:818bcc6723561f207e5b5453efe9703f34bc8781d11ce9b8be286bb415eb4962", size = 678335, upload-time = "2025-11-14T09:32:20.107Z" }, + { url = "https://files.pythonhosted.org/packages/f4/d2/29e5e536adc07bc3d33dd09f3f7cf844bf7b4981820dc2a91dd810f3c782/pyobjc_core-12.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:01c0cf500596f03e21c23aef9b5f326b9fb1f8f118cf0d8b66749b6cf4cbb37a", size = 677370, upload-time = "2025-11-14T09:33:05.273Z" }, + { url = "https://files.pythonhosted.org/packages/1b/f0/4b4ed8924cd04e425f2a07269943018d43949afad1c348c3ed4d9d032787/pyobjc_core-12.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:177aaca84bb369a483e4961186704f64b2697708046745f8167e818d968c88fc", size = 719586, upload-time = "2025-11-14T09:33:53.302Z" }, +] + +[[package]] +name = "pyobjc-framework-applicationservices" +version = "12.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyobjc-core", marker = "sys_platform == 'darwin'" }, + { name = "pyobjc-framework-cocoa", marker = "sys_platform == 'darwin'" }, + { name = "pyobjc-framework-coretext", marker = "sys_platform == 'darwin'" }, + { name = "pyobjc-framework-quartz", marker = "sys_platform == 'darwin'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/be/6a/d4e613c8e926a5744fc47a9e9fea08384a510dc4f27d844f7ad7a2d793bd/pyobjc_framework_applicationservices-12.1.tar.gz", hash = "sha256:c06abb74f119bc27aeb41bf1aef8102c0ae1288aec1ac8665ea186a067a8945b", size = 103247, upload-time = "2025-11-14T10:08:52.18Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/37/a7/55fa88def5c02732c4b747606ff1cbce6e1f890734bbd00f5596b21eaa02/pyobjc_framework_applicationservices-12.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:c8f6e2fb3b3e9214ab4864ef04eee18f592b46a986c86ea0113448b310520532", size = 32835, upload-time = "2025-11-14T09:36:11.855Z" }, + { url = "https://files.pythonhosted.org/packages/fc/21/79e42ee836f1010f5fe9e97d2817a006736bd287c15a3674c399190a2e77/pyobjc_framework_applicationservices-12.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:bd1f4dbb38234a24ae6819f5e22485cf7dd3dd4074ff3bf9a9fdb4c01a3b4a38", size = 32859, upload-time = "2025-11-14T09:36:15.208Z" }, + { url = "https://files.pythonhosted.org/packages/66/3a/0f1d4dcf2345e875e5ea9761d5a70969e241d24089133d21f008dde596f5/pyobjc_framework_applicationservices-12.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:8a5d2845249b6a85ba9e320a9848468c3f8cd6f59605a9a43f406a7810eaa830", size = 33115, upload-time = "2025-11-14T09:36:18.384Z" }, +] + +[[package]] +name = "pyobjc-framework-cocoa" +version = "12.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyobjc-core", marker = "sys_platform == 'darwin'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/02/a3/16ca9a15e77c061a9250afbae2eae26f2e1579eb8ca9462ae2d2c71e1169/pyobjc_framework_cocoa-12.1.tar.gz", hash = "sha256:5556c87db95711b985d5efdaaf01c917ddd41d148b1e52a0c66b1a2e2c5c1640", size = 2772191, upload-time = "2025-11-14T10:13:02.069Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/95/bf/ee4f27ec3920d5c6fc63c63e797c5b2cc4e20fe439217085d01ea5b63856/pyobjc_framework_cocoa-12.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:547c182837214b7ec4796dac5aee3aa25abc665757b75d7f44f83c994bcb0858", size = 384590, upload-time = "2025-11-14T09:41:17.336Z" }, + { url = "https://files.pythonhosted.org/packages/ad/31/0c2e734165abb46215797bd830c4bdcb780b699854b15f2b6240515edcc6/pyobjc_framework_cocoa-12.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:5a3dcd491cacc2f5a197142b3c556d8aafa3963011110102a093349017705118", size = 384689, upload-time = "2025-11-14T09:41:41.478Z" }, + { url = "https://files.pythonhosted.org/packages/23/3b/b9f61be7b9f9b4e0a6db18b3c35c4c4d589f2d04e963e2174d38c6555a92/pyobjc_framework_cocoa-12.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:914b74328c22d8ca261d78c23ef2befc29776e0b85555973927b338c5734ca44", size = 388843, upload-time = "2025-11-14T09:42:05.719Z" }, +] + +[[package]] +name = "pyobjc-framework-coretext" +version = "12.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyobjc-core", marker = "sys_platform == 'darwin'" }, + { name = "pyobjc-framework-cocoa", marker = "sys_platform == 'darwin'" }, + { name = "pyobjc-framework-quartz", marker = "sys_platform == 'darwin'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/29/da/682c9c92a39f713bd3c56e7375fa8f1b10ad558ecb075258ab6f1cdd4a6d/pyobjc_framework_coretext-12.1.tar.gz", hash = "sha256:e0adb717738fae395dc645c9e8a10bb5f6a4277e73cba8fa2a57f3b518e71da5", size = 90124, upload-time = "2025-11-14T10:14:38.596Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cd/0f/ddf45bf0e3ba4fbdc7772de4728fd97ffc34a0b5a15e1ab1115b202fe4ae/pyobjc_framework_coretext-12.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:d246fa654bdbf43bae3969887d58f0b336c29b795ad55a54eb76397d0e62b93c", size = 30108, upload-time = "2025-11-14T09:47:04.228Z" }, + { url = "https://files.pythonhosted.org/packages/20/a2/a3974e3e807c68e23a9d7db66fc38ac54f7ecd2b7a9237042006699a76e1/pyobjc_framework_coretext-12.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:7cbb2c28580e6704ce10b9a991ccd9563a22b3a75f67c36cf612544bd8b21b5f", size = 30110, upload-time = "2025-11-14T09:47:07.518Z" }, + { url = "https://files.pythonhosted.org/packages/0f/5d/85e059349e9cfbd57269a1f11f56747b3ff5799a3bcbd95485f363c623d8/pyobjc_framework_coretext-12.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:14100d1e39efb30f57869671fb6fce8d668f80c82e25e7930fb364866e5c0dab", size = 30697, upload-time = "2025-11-14T09:47:10.932Z" }, +] + +[[package]] +name = "pyobjc-framework-quartz" +version = "12.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyobjc-core", marker = "sys_platform == 'darwin'" }, + { name = "pyobjc-framework-cocoa", marker = "sys_platform == 'darwin'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/94/18/cc59f3d4355c9456fc945eae7fe8797003c4da99212dd531ad1b0de8a0c6/pyobjc_framework_quartz-12.1.tar.gz", hash = "sha256:27f782f3513ac88ec9b6c82d9767eef95a5cf4175ce88a1e5a65875fee799608", size = 3159099, upload-time = "2025-11-14T10:21:24.31Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e9/9b/780f057e5962f690f23fdff1083a4cfda5a96d5b4d3bb49505cac4f624f2/pyobjc_framework_quartz-12.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:7730cdce46c7e985535b5a42c31381af4aa6556e5642dc55b5e6597595e57a16", size = 218798, upload-time = "2025-11-14T10:00:01.236Z" }, + { url = "https://files.pythonhosted.org/packages/ba/2d/e8f495328101898c16c32ac10e7b14b08ff2c443a756a76fd1271915f097/pyobjc_framework_quartz-12.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:629b7971b1b43a11617f1460cd218bd308dfea247cd4ee3842eb40ca6f588860", size = 219206, upload-time = "2025-11-14T10:00:15.623Z" }, + { url = "https://files.pythonhosted.org/packages/67/43/b1f0ad3b842ab150a7e6b7d97f6257eab6af241b4c7d14cb8e7fde9214b8/pyobjc_framework_quartz-12.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:53b84e880c358ba1ddcd7e8d5ea0407d760eca58b96f0d344829162cda5f37b3", size = 224317, upload-time = "2025-11-14T10:00:30.703Z" }, +] + +[[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 = "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 = "python-dotenv" +version = "1.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/82/ed/0301aeeac3e5353ef3d94b6ec08bbcabd04a72018415dcb29e588514bba8/python_dotenv-1.2.2.tar.gz", hash = "sha256:2c371a91fbd7ba082c2c1dc1f8bf89ca22564a087c2c287cd9b662adde799cf3", size = 50135, upload-time = "2026-03-01T16:00:26.196Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0b/d7/1959b9648791274998a9c3526f6d0ec8fd2233e4d4acce81bbae76b44b2a/python_dotenv-1.2.2-py3-none-any.whl", hash = "sha256:1d8214789a24de455a8b8bd8ae6fe3c6b69a5e3d64aa8a8e5d68e694bbcb285a", size = 22101, upload-time = "2026-03-01T16:00:25.09Z" }, +] + +[[package]] +name = "python-xlib" +version = "0.33" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "six" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/86/f5/8c0653e5bb54e0cbdfe27bf32d41f27bc4e12faa8742778c17f2a71be2c0/python-xlib-0.33.tar.gz", hash = "sha256:55af7906a2c75ce6cb280a584776080602444f75815a7aff4d287bb2d7018b32", size = 269068, upload-time = "2022-12-25T18:53:00.824Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fc/b8/ff33610932e0ee81ae7f1269c890f697d56ff74b9f5b2ee5d9b7fa2c5355/python_xlib-0.33-py2.py3-none-any.whl", hash = "sha256:c3534038d42e0df2f1392a1b30a15a4ff5fdc2b86cfa94f072bf11b10a164398", size = 182185, upload-time = "2022-12-25T18:52:58.662Z" }, +] + +[[package]] +name = "pyvers" +version = "0.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/32/99/23c73a1298b1c642d8ebdd78e1db4daf1e474152e6839df4f5c93357a3db/pyvers-0.2.2.tar.gz", hash = "sha256:205026bcd0b4c09198cb3a32f243fd179ef012882ce16d93dcb755320acd56f7", size = 12104, upload-time = "2026-01-23T14:12:07.619Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/36/bf/ea840f706b7824dd57220484465995309c8c217995ddb7ce4b262240e912/pyvers-0.2.2-py3-none-any.whl", hash = "sha256:c4696408a0b15fbaa90df33d3bc579cf23a74a73541858f5470216f12f51f3b1", size = 11569, upload-time = "2026-01-23T14:12:06.246Z" }, +] + +[[package]] +name = "pyyaml" +version = "6.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/05/8e/961c0007c59b8dd7729d542c61a4d537767a59645b82a0b521206e1e25c2/pyyaml-6.0.3.tar.gz", hash = "sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f", size = 130960, upload-time = "2025-09-25T21:33:16.546Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/33/422b98d2195232ca1826284a76852ad5a86fe23e31b009c9886b2d0fb8b2/pyyaml-6.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7f047e29dcae44602496db43be01ad42fc6f1cc0d8cd6c83d342306c32270196", size = 182063, upload-time = "2025-09-25T21:32:11.445Z" }, + { url = "https://files.pythonhosted.org/packages/89/a0/6cf41a19a1f2f3feab0e9c0b74134aa2ce6849093d5517a0c550fe37a648/pyyaml-6.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:fc09d0aa354569bc501d4e787133afc08552722d3ab34836a80547331bb5d4a0", size = 173973, upload-time = "2025-09-25T21:32:12.492Z" }, + { url = "https://files.pythonhosted.org/packages/ed/23/7a778b6bd0b9a8039df8b1b1d80e2e2ad78aa04171592c8a5c43a56a6af4/pyyaml-6.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9149cad251584d5fb4981be1ecde53a1ca46c891a79788c0df828d2f166bda28", size = 775116, upload-time = "2025-09-25T21:32:13.652Z" }, + { url = "https://files.pythonhosted.org/packages/65/30/d7353c338e12baef4ecc1b09e877c1970bd3382789c159b4f89d6a70dc09/pyyaml-6.0.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5fdec68f91a0c6739b380c83b951e2c72ac0197ace422360e6d5a959d8d97b2c", size = 844011, upload-time = "2025-09-25T21:32:15.21Z" }, + { url = "https://files.pythonhosted.org/packages/8b/9d/b3589d3877982d4f2329302ef98a8026e7f4443c765c46cfecc8858c6b4b/pyyaml-6.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ba1cc08a7ccde2d2ec775841541641e4548226580ab850948cbfda66a1befcdc", size = 807870, upload-time = "2025-09-25T21:32:16.431Z" }, + { url = "https://files.pythonhosted.org/packages/05/c0/b3be26a015601b822b97d9149ff8cb5ead58c66f981e04fedf4e762f4bd4/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8dc52c23056b9ddd46818a57b78404882310fb473d63f17b07d5c40421e47f8e", size = 761089, upload-time = "2025-09-25T21:32:17.56Z" }, + { url = "https://files.pythonhosted.org/packages/be/8e/98435a21d1d4b46590d5459a22d88128103f8da4c2d4cb8f14f2a96504e1/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:41715c910c881bc081f1e8872880d3c650acf13dfa8214bad49ed4cede7c34ea", size = 790181, upload-time = "2025-09-25T21:32:18.834Z" }, + { url = "https://files.pythonhosted.org/packages/74/93/7baea19427dcfbe1e5a372d81473250b379f04b1bd3c4c5ff825e2327202/pyyaml-6.0.3-cp312-cp312-win32.whl", hash = "sha256:96b533f0e99f6579b3d4d4995707cf36df9100d67e0c8303a0c55b27b5f99bc5", size = 137658, upload-time = "2025-09-25T21:32:20.209Z" }, + { url = "https://files.pythonhosted.org/packages/86/bf/899e81e4cce32febab4fb42bb97dcdf66bc135272882d1987881a4b519e9/pyyaml-6.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:5fcd34e47f6e0b794d17de1b4ff496c00986e1c83f7ab2fb8fcfe9616ff7477b", size = 154003, upload-time = "2025-09-25T21:32:21.167Z" }, + { url = "https://files.pythonhosted.org/packages/1a/08/67bd04656199bbb51dbed1439b7f27601dfb576fb864099c7ef0c3e55531/pyyaml-6.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:64386e5e707d03a7e172c0701abfb7e10f0fb753ee1d773128192742712a98fd", size = 140344, upload-time = "2025-09-25T21:32:22.617Z" }, + { url = "https://files.pythonhosted.org/packages/d1/11/0fd08f8192109f7169db964b5707a2f1e8b745d4e239b784a5a1dd80d1db/pyyaml-6.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:8da9669d359f02c0b91ccc01cac4a67f16afec0dac22c2ad09f46bee0697eba8", size = 181669, upload-time = "2025-09-25T21:32:23.673Z" }, + { url = "https://files.pythonhosted.org/packages/b1/16/95309993f1d3748cd644e02e38b75d50cbc0d9561d21f390a76242ce073f/pyyaml-6.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:2283a07e2c21a2aa78d9c4442724ec1eb15f5e42a723b99cb3d822d48f5f7ad1", size = 173252, upload-time = "2025-09-25T21:32:25.149Z" }, + { url = "https://files.pythonhosted.org/packages/50/31/b20f376d3f810b9b2371e72ef5adb33879b25edb7a6d072cb7ca0c486398/pyyaml-6.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ee2922902c45ae8ccada2c5b501ab86c36525b883eff4255313a253a3160861c", size = 767081, upload-time = "2025-09-25T21:32:26.575Z" }, + { url = "https://files.pythonhosted.org/packages/49/1e/a55ca81e949270d5d4432fbbd19dfea5321eda7c41a849d443dc92fd1ff7/pyyaml-6.0.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a33284e20b78bd4a18c8c2282d549d10bc8408a2a7ff57653c0cf0b9be0afce5", size = 841159, upload-time = "2025-09-25T21:32:27.727Z" }, + { url = "https://files.pythonhosted.org/packages/74/27/e5b8f34d02d9995b80abcef563ea1f8b56d20134d8f4e5e81733b1feceb2/pyyaml-6.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0f29edc409a6392443abf94b9cf89ce99889a1dd5376d94316ae5145dfedd5d6", size = 801626, upload-time = "2025-09-25T21:32:28.878Z" }, + { url = "https://files.pythonhosted.org/packages/f9/11/ba845c23988798f40e52ba45f34849aa8a1f2d4af4b798588010792ebad6/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f7057c9a337546edc7973c0d3ba84ddcdf0daa14533c2065749c9075001090e6", size = 753613, upload-time = "2025-09-25T21:32:30.178Z" }, + { url = "https://files.pythonhosted.org/packages/3d/e0/7966e1a7bfc0a45bf0a7fb6b98ea03fc9b8d84fa7f2229e9659680b69ee3/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:eda16858a3cab07b80edaf74336ece1f986ba330fdb8ee0d6c0d68fe82bc96be", size = 794115, upload-time = "2025-09-25T21:32:31.353Z" }, + { url = "https://files.pythonhosted.org/packages/de/94/980b50a6531b3019e45ddeada0626d45fa85cbe22300844a7983285bed3b/pyyaml-6.0.3-cp313-cp313-win32.whl", hash = "sha256:d0eae10f8159e8fdad514efdc92d74fd8d682c933a6dd088030f3834bc8e6b26", size = 137427, upload-time = "2025-09-25T21:32:32.58Z" }, + { url = "https://files.pythonhosted.org/packages/97/c9/39d5b874e8b28845e4ec2202b5da735d0199dbe5b8fb85f91398814a9a46/pyyaml-6.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:79005a0d97d5ddabfeeea4cf676af11e647e41d81c9a7722a193022accdb6b7c", size = 154090, upload-time = "2025-09-25T21:32:33.659Z" }, + { url = "https://files.pythonhosted.org/packages/73/e8/2bdf3ca2090f68bb3d75b44da7bbc71843b19c9f2b9cb9b0f4ab7a5a4329/pyyaml-6.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:5498cd1645aa724a7c71c8f378eb29ebe23da2fc0d7a08071d89469bf1d2defb", size = 140246, upload-time = "2025-09-25T21:32:34.663Z" }, +] + +[[package]] +name = "referencing" +version = "0.37.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, + { name = "rpds-py" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/22/f5/df4e9027acead3ecc63e50fe1e36aca1523e1719559c499951bb4b53188f/referencing-0.37.0.tar.gz", hash = "sha256:44aefc3142c5b842538163acb373e24cce6632bd54bdb01b21ad5863489f50d8", size = 78036, upload-time = "2025-10-13T15:30:48.871Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl", hash = "sha256:381329a9f99628c9069361716891d34ad94af76e461dcb0335825aecc7692231", size = 26766, upload-time = "2025-10-13T15:30:47.625Z" }, +] + +[[package]] +name = "requests" +version = "2.33.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "charset-normalizer" }, + { name = "idna" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5f/a4/98b9c7c6428a668bf7e42ebb7c79d576a1c3c1e3ae2d47e674b468388871/requests-2.33.1.tar.gz", hash = "sha256:18817f8c57c6263968bc123d237e3b8b08ac046f5456bd1e307ee8f4250d3517", size = 134120, upload-time = "2026-03-30T16:09:15.531Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d7/8e/7540e8a2036f79a125c1d2ebadf69ed7901608859186c856fa0388ef4197/requests-2.33.1-py3-none-any.whl", hash = "sha256:4e6d1ef462f3626a1f0a0a9c42dd93c63bad33f9f1c1937509b8c5c8718ab56a", size = 64947, upload-time = "2026-03-30T16:09:13.83Z" }, +] + +[[package]] +name = "rerun-sdk" +version = "0.31.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, + { name = "numpy" }, + { name = "pillow" }, + { name = "pyarrow" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/6d/eb/48ff09062919646bed150e76313debc7b815c23de12dbb999332fb0bddc2/rerun_sdk-0.31.3-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:d3ed6c246edf6f809df5fbdaeafbb335d1295930b73d9d53b17e7d9fb1f5e8bc", size = 121398857, upload-time = "2026-04-14T10:10:24.611Z" }, + { url = "https://files.pythonhosted.org/packages/e2/8c/219d612b61d2fde3731068cd3ba1ef1e9bb646e06923361574b52529fa96/rerun_sdk-0.31.3-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:90deef848af60d544071e219ce2f608f64090a50b75d0cd951c24933a0260edd", size = 130846353, upload-time = "2026-04-14T10:10:31.311Z" }, + { url = "https://files.pythonhosted.org/packages/f7/f4/d4350829d5bd74005ff08b194a21c24a570c9db67d90c14ca29cf0b063bf/rerun_sdk-0.31.3-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:b285dda76f4bf01d617204c09bab8543cddd224cc8fe159f3bf38458bf398998", size = 135093900, upload-time = "2026-04-14T10:10:36.404Z" }, + { url = "https://files.pythonhosted.org/packages/5d/99/f1d5ba615a282ae18919c1c3c1975cbd14c66032ed3c574ac0153fbbb567/rerun_sdk-0.31.3-cp310-abi3-win_amd64.whl", hash = "sha256:427b90d5157189f134b0e43d3378c83f548ac858cc0b81553ecb6df4b7977a40", size = 116476172, upload-time = "2026-04-14T10:10:41.209Z" }, +] + +[[package]] +name = "rich" +version = "14.3.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e9/67/cae617f1351490c25a4b8ac3b8b63a4dda609295d8222bad12242dfdc629/rich-14.3.4.tar.gz", hash = "sha256:817e02727f2b25b40ef56f5aa2217f400c8489f79ca8f46ea2b70dd5e14558a9", size = 230524, upload-time = "2026-04-11T02:57:45.419Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/76/6d163cfac87b632216f71879e6b2cf17163f773ff59c00b5ff4900a80fa3/rich-14.3.4-py3-none-any.whl", hash = "sha256:07e7adb4690f68864777b1450859253bed81a99a31ac321ac1817b2313558952", size = 310480, upload-time = "2026-04-11T02:57:47.484Z" }, +] + +[[package]] +name = "rpds-py" +version = "0.30.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/20/af/3f2f423103f1113b36230496629986e0ef7e199d2aa8392452b484b38ced/rpds_py-0.30.0.tar.gz", hash = "sha256:dd8ff7cf90014af0c0f787eea34794ebf6415242ee1d6fa91eaba725cc441e84", size = 69469, upload-time = "2025-11-30T20:24:38.837Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/03/e7/98a2f4ac921d82f33e03f3835f5bf3a4a40aa1bfdc57975e74a97b2b4bdd/rpds_py-0.30.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:a161f20d9a43006833cd7068375a94d035714d73a172b681d8881820600abfad", size = 375086, upload-time = "2025-11-30T20:22:17.93Z" }, + { url = "https://files.pythonhosted.org/packages/4d/a1/bca7fd3d452b272e13335db8d6b0b3ecde0f90ad6f16f3328c6fb150c889/rpds_py-0.30.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6abc8880d9d036ecaafe709079969f56e876fcf107f7a8e9920ba6d5a3878d05", size = 359053, upload-time = "2025-11-30T20:22:19.297Z" }, + { url = "https://files.pythonhosted.org/packages/65/1c/ae157e83a6357eceff62ba7e52113e3ec4834a84cfe07fa4b0757a7d105f/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ca28829ae5f5d569bb62a79512c842a03a12576375d5ece7d2cadf8abe96ec28", size = 390763, upload-time = "2025-11-30T20:22:21.661Z" }, + { url = "https://files.pythonhosted.org/packages/d4/36/eb2eb8515e2ad24c0bd43c3ee9cd74c33f7ca6430755ccdb240fd3144c44/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:a1010ed9524c73b94d15919ca4d41d8780980e1765babf85f9a2f90d247153dd", size = 408951, upload-time = "2025-11-30T20:22:23.408Z" }, + { url = "https://files.pythonhosted.org/packages/d6/65/ad8dc1784a331fabbd740ef6f71ce2198c7ed0890dab595adb9ea2d775a1/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f8d1736cfb49381ba528cd5baa46f82fdc65c06e843dab24dd70b63d09121b3f", size = 514622, upload-time = "2025-11-30T20:22:25.16Z" }, + { url = "https://files.pythonhosted.org/packages/63/8e/0cfa7ae158e15e143fe03993b5bcd743a59f541f5952e1546b1ac1b5fd45/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d948b135c4693daff7bc2dcfc4ec57237a29bd37e60c2fabf5aff2bbacf3e2f1", size = 414492, upload-time = "2025-11-30T20:22:26.505Z" }, + { url = "https://files.pythonhosted.org/packages/60/1b/6f8f29f3f995c7ffdde46a626ddccd7c63aefc0efae881dc13b6e5d5bb16/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:47f236970bccb2233267d89173d3ad2703cd36a0e2a6e92d0560d333871a3d23", size = 394080, upload-time = "2025-11-30T20:22:27.934Z" }, + { url = "https://files.pythonhosted.org/packages/6d/d5/a266341051a7a3ca2f4b750a3aa4abc986378431fc2da508c5034d081b70/rpds_py-0.30.0-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:2e6ecb5a5bcacf59c3f912155044479af1d0b6681280048b338b28e364aca1f6", size = 408680, upload-time = "2025-11-30T20:22:29.341Z" }, + { url = "https://files.pythonhosted.org/packages/10/3b/71b725851df9ab7a7a4e33cf36d241933da66040d195a84781f49c50490c/rpds_py-0.30.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:a8fa71a2e078c527c3e9dc9fc5a98c9db40bcc8a92b4e8858e36d329f8684b51", size = 423589, upload-time = "2025-11-30T20:22:31.469Z" }, + { url = "https://files.pythonhosted.org/packages/00/2b/e59e58c544dc9bd8bd8384ecdb8ea91f6727f0e37a7131baeff8d6f51661/rpds_py-0.30.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:73c67f2db7bc334e518d097c6d1e6fed021bbc9b7d678d6cc433478365d1d5f5", size = 573289, upload-time = "2025-11-30T20:22:32.997Z" }, + { url = "https://files.pythonhosted.org/packages/da/3e/a18e6f5b460893172a7d6a680e86d3b6bc87a54c1f0b03446a3c8c7b588f/rpds_py-0.30.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:5ba103fb455be00f3b1c2076c9d4264bfcb037c976167a6047ed82f23153f02e", size = 599737, upload-time = "2025-11-30T20:22:34.419Z" }, + { url = "https://files.pythonhosted.org/packages/5c/e2/714694e4b87b85a18e2c243614974413c60aa107fd815b8cbc42b873d1d7/rpds_py-0.30.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:7cee9c752c0364588353e627da8a7e808a66873672bcb5f52890c33fd965b394", size = 563120, upload-time = "2025-11-30T20:22:35.903Z" }, + { url = "https://files.pythonhosted.org/packages/6f/ab/d5d5e3bcedb0a77f4f613706b750e50a5a3ba1c15ccd3665ecc636c968fd/rpds_py-0.30.0-cp312-cp312-win32.whl", hash = "sha256:1ab5b83dbcf55acc8b08fc62b796ef672c457b17dbd7820a11d6c52c06839bdf", size = 223782, upload-time = "2025-11-30T20:22:37.271Z" }, + { url = "https://files.pythonhosted.org/packages/39/3b/f786af9957306fdc38a74cef405b7b93180f481fb48453a114bb6465744a/rpds_py-0.30.0-cp312-cp312-win_amd64.whl", hash = "sha256:a090322ca841abd453d43456ac34db46e8b05fd9b3b4ac0c78bcde8b089f959b", size = 240463, upload-time = "2025-11-30T20:22:39.021Z" }, + { url = "https://files.pythonhosted.org/packages/f3/d2/b91dc748126c1559042cfe41990deb92c4ee3e2b415f6b5234969ffaf0cc/rpds_py-0.30.0-cp312-cp312-win_arm64.whl", hash = "sha256:669b1805bd639dd2989b281be2cfd951c6121b65e729d9b843e9639ef1fd555e", size = 230868, upload-time = "2025-11-30T20:22:40.493Z" }, + { url = "https://files.pythonhosted.org/packages/ed/dc/d61221eb88ff410de3c49143407f6f3147acf2538c86f2ab7ce65ae7d5f9/rpds_py-0.30.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:f83424d738204d9770830d35290ff3273fbb02b41f919870479fab14b9d303b2", size = 374887, upload-time = "2025-11-30T20:22:41.812Z" }, + { url = "https://files.pythonhosted.org/packages/fd/32/55fb50ae104061dbc564ef15cc43c013dc4a9f4527a1f4d99baddf56fe5f/rpds_py-0.30.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:e7536cd91353c5273434b4e003cbda89034d67e7710eab8761fd918ec6c69cf8", size = 358904, upload-time = "2025-11-30T20:22:43.479Z" }, + { url = "https://files.pythonhosted.org/packages/58/70/faed8186300e3b9bdd138d0273109784eea2396c68458ed580f885dfe7ad/rpds_py-0.30.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2771c6c15973347f50fece41fc447c054b7ac2ae0502388ce3b6738cd366e3d4", size = 389945, upload-time = "2025-11-30T20:22:44.819Z" }, + { url = "https://files.pythonhosted.org/packages/bd/a8/073cac3ed2c6387df38f71296d002ab43496a96b92c823e76f46b8af0543/rpds_py-0.30.0-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:0a59119fc6e3f460315fe9d08149f8102aa322299deaa5cab5b40092345c2136", size = 407783, upload-time = "2025-11-30T20:22:46.103Z" }, + { url = "https://files.pythonhosted.org/packages/77/57/5999eb8c58671f1c11eba084115e77a8899d6e694d2a18f69f0ba471ec8b/rpds_py-0.30.0-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:76fec018282b4ead0364022e3c54b60bf368b9d926877957a8624b58419169b7", size = 515021, upload-time = "2025-11-30T20:22:47.458Z" }, + { url = "https://files.pythonhosted.org/packages/e0/af/5ab4833eadc36c0a8ed2bc5c0de0493c04f6c06de223170bd0798ff98ced/rpds_py-0.30.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:692bef75a5525db97318e8cd061542b5a79812d711ea03dbc1f6f8dbb0c5f0d2", size = 414589, upload-time = "2025-11-30T20:22:48.872Z" }, + { url = "https://files.pythonhosted.org/packages/b7/de/f7192e12b21b9e9a68a6d0f249b4af3fdcdff8418be0767a627564afa1f1/rpds_py-0.30.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9027da1ce107104c50c81383cae773ef5c24d296dd11c99e2629dbd7967a20c6", size = 394025, upload-time = "2025-11-30T20:22:50.196Z" }, + { url = "https://files.pythonhosted.org/packages/91/c4/fc70cd0249496493500e7cc2de87504f5aa6509de1e88623431fec76d4b6/rpds_py-0.30.0-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:9cf69cdda1f5968a30a359aba2f7f9aa648a9ce4b580d6826437f2b291cfc86e", size = 408895, upload-time = "2025-11-30T20:22:51.87Z" }, + { url = "https://files.pythonhosted.org/packages/58/95/d9275b05ab96556fefff73a385813eb66032e4c99f411d0795372d9abcea/rpds_py-0.30.0-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:a4796a717bf12b9da9d3ad002519a86063dcac8988b030e405704ef7d74d2d9d", size = 422799, upload-time = "2025-11-30T20:22:53.341Z" }, + { url = "https://files.pythonhosted.org/packages/06/c1/3088fc04b6624eb12a57eb814f0d4997a44b0d208d6cace713033ff1a6ba/rpds_py-0.30.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:5d4c2aa7c50ad4728a094ebd5eb46c452e9cb7edbfdb18f9e1221f597a73e1e7", size = 572731, upload-time = "2025-11-30T20:22:54.778Z" }, + { url = "https://files.pythonhosted.org/packages/d8/42/c612a833183b39774e8ac8fecae81263a68b9583ee343db33ab571a7ce55/rpds_py-0.30.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:ba81a9203d07805435eb06f536d95a266c21e5b2dfbf6517748ca40c98d19e31", size = 599027, upload-time = "2025-11-30T20:22:56.212Z" }, + { url = "https://files.pythonhosted.org/packages/5f/60/525a50f45b01d70005403ae0e25f43c0384369ad24ffe46e8d9068b50086/rpds_py-0.30.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:945dccface01af02675628334f7cf49c2af4c1c904748efc5cf7bbdf0b579f95", size = 563020, upload-time = "2025-11-30T20:22:58.2Z" }, + { url = "https://files.pythonhosted.org/packages/0b/5d/47c4655e9bcd5ca907148535c10e7d489044243cc9941c16ed7cd53be91d/rpds_py-0.30.0-cp313-cp313-win32.whl", hash = "sha256:b40fb160a2db369a194cb27943582b38f79fc4887291417685f3ad693c5a1d5d", size = 223139, upload-time = "2025-11-30T20:23:00.209Z" }, + { url = "https://files.pythonhosted.org/packages/f2/e1/485132437d20aa4d3e1d8b3fb5a5e65aa8139f1e097080c2a8443201742c/rpds_py-0.30.0-cp313-cp313-win_amd64.whl", hash = "sha256:806f36b1b605e2d6a72716f321f20036b9489d29c51c91f4dd29a3e3afb73b15", size = 240224, upload-time = "2025-11-30T20:23:02.008Z" }, + { url = "https://files.pythonhosted.org/packages/24/95/ffd128ed1146a153d928617b0ef673960130be0009c77d8fbf0abe306713/rpds_py-0.30.0-cp313-cp313-win_arm64.whl", hash = "sha256:d96c2086587c7c30d44f31f42eae4eac89b60dabbac18c7669be3700f13c3ce1", size = 230645, upload-time = "2025-11-30T20:23:03.43Z" }, + { url = "https://files.pythonhosted.org/packages/ff/1b/b10de890a0def2a319a2626334a7f0ae388215eb60914dbac8a3bae54435/rpds_py-0.30.0-cp313-cp313t-macosx_10_12_x86_64.whl", hash = "sha256:eb0b93f2e5c2189ee831ee43f156ed34e2a89a78a66b98cadad955972548be5a", size = 364443, upload-time = "2025-11-30T20:23:04.878Z" }, + { url = "https://files.pythonhosted.org/packages/0d/bf/27e39f5971dc4f305a4fb9c672ca06f290f7c4e261c568f3dea16a410d47/rpds_py-0.30.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:922e10f31f303c7c920da8981051ff6d8c1a56207dbdf330d9047f6d30b70e5e", size = 353375, upload-time = "2025-11-30T20:23:06.342Z" }, + { url = "https://files.pythonhosted.org/packages/40/58/442ada3bba6e8e6615fc00483135c14a7538d2ffac30e2d933ccf6852232/rpds_py-0.30.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cdc62c8286ba9bf7f47befdcea13ea0e26bf294bda99758fd90535cbaf408000", size = 383850, upload-time = "2025-11-30T20:23:07.825Z" }, + { url = "https://files.pythonhosted.org/packages/14/14/f59b0127409a33c6ef6f5c1ebd5ad8e32d7861c9c7adfa9a624fc3889f6c/rpds_py-0.30.0-cp313-cp313t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:47f9a91efc418b54fb8190a6b4aa7813a23fb79c51f4bb84e418f5476c38b8db", size = 392812, upload-time = "2025-11-30T20:23:09.228Z" }, + { url = "https://files.pythonhosted.org/packages/b3/66/e0be3e162ac299b3a22527e8913767d869e6cc75c46bd844aa43fb81ab62/rpds_py-0.30.0-cp313-cp313t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1f3587eb9b17f3789ad50824084fa6f81921bbf9a795826570bda82cb3ed91f2", size = 517841, upload-time = "2025-11-30T20:23:11.186Z" }, + { url = "https://files.pythonhosted.org/packages/3d/55/fa3b9cf31d0c963ecf1ba777f7cf4b2a2c976795ac430d24a1f43d25a6ba/rpds_py-0.30.0-cp313-cp313t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:39c02563fc592411c2c61d26b6c5fe1e51eaa44a75aa2c8735ca88b0d9599daa", size = 408149, upload-time = "2025-11-30T20:23:12.864Z" }, + { url = "https://files.pythonhosted.org/packages/60/ca/780cf3b1a32b18c0f05c441958d3758f02544f1d613abf9488cd78876378/rpds_py-0.30.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:51a1234d8febafdfd33a42d97da7a43f5dcb120c1060e352a3fbc0c6d36e2083", size = 383843, upload-time = "2025-11-30T20:23:14.638Z" }, + { url = "https://files.pythonhosted.org/packages/82/86/d5f2e04f2aa6247c613da0c1dd87fcd08fa17107e858193566048a1e2f0a/rpds_py-0.30.0-cp313-cp313t-manylinux_2_31_riscv64.whl", hash = "sha256:eb2c4071ab598733724c08221091e8d80e89064cd472819285a9ab0f24bcedb9", size = 396507, upload-time = "2025-11-30T20:23:16.105Z" }, + { url = "https://files.pythonhosted.org/packages/4b/9a/453255d2f769fe44e07ea9785c8347edaf867f7026872e76c1ad9f7bed92/rpds_py-0.30.0-cp313-cp313t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:6bdfdb946967d816e6adf9a3d8201bfad269c67efe6cefd7093ef959683c8de0", size = 414949, upload-time = "2025-11-30T20:23:17.539Z" }, + { url = "https://files.pythonhosted.org/packages/a3/31/622a86cdc0c45d6df0e9ccb6becdba5074735e7033c20e401a6d9d0e2ca0/rpds_py-0.30.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:c77afbd5f5250bf27bf516c7c4a016813eb2d3e116139aed0096940c5982da94", size = 565790, upload-time = "2025-11-30T20:23:19.029Z" }, + { url = "https://files.pythonhosted.org/packages/1c/5d/15bbf0fb4a3f58a3b1c67855ec1efcc4ceaef4e86644665fff03e1b66d8d/rpds_py-0.30.0-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:61046904275472a76c8c90c9ccee9013d70a6d0f73eecefd38c1ae7c39045a08", size = 590217, upload-time = "2025-11-30T20:23:20.885Z" }, + { url = "https://files.pythonhosted.org/packages/6d/61/21b8c41f68e60c8cc3b2e25644f0e3681926020f11d06ab0b78e3c6bbff1/rpds_py-0.30.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:4c5f36a861bc4b7da6516dbdf302c55313afa09b81931e8280361a4f6c9a2d27", size = 555806, upload-time = "2025-11-30T20:23:22.488Z" }, + { url = "https://files.pythonhosted.org/packages/f9/39/7e067bb06c31de48de3eb200f9fc7c58982a4d3db44b07e73963e10d3be9/rpds_py-0.30.0-cp313-cp313t-win32.whl", hash = "sha256:3d4a69de7a3e50ffc214ae16d79d8fbb0922972da0356dcf4d0fdca2878559c6", size = 211341, upload-time = "2025-11-30T20:23:24.449Z" }, + { url = "https://files.pythonhosted.org/packages/0a/4d/222ef0b46443cf4cf46764d9c630f3fe4abaa7245be9417e56e9f52b8f65/rpds_py-0.30.0-cp313-cp313t-win_amd64.whl", hash = "sha256:f14fc5df50a716f7ece6a80b6c78bb35ea2ca47c499e422aa4463455dd96d56d", size = 225768, upload-time = "2025-11-30T20:23:25.908Z" }, +] + +[[package]] +name = "rsl-rl-lib" +version = "3.1.0" +source = { directory = "rsl_rl" } +dependencies = [ + { name = "gitpython" }, + { name = "numpy" }, + { name = "onnx" }, + { name = "tensordict" }, + { name = "torch" }, + { name = "torchvision" }, +] + +[package.metadata] +requires-dist = [ + { name = "gitpython" }, + { name = "numpy" }, + { name = "onnx" }, + { name = "tensordict" }, + { name = "torch" }, + { name = "torchvision" }, +] + +[[package]] +name = "rtree" +version = "1.4.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/95/09/7302695875a019514de9a5dd17b8320e7a19d6e7bc8f85dcfb79a4ce2da3/rtree-1.4.1.tar.gz", hash = "sha256:c6b1b3550881e57ebe530cc6cffefc87cd9bf49c30b37b894065a9f810875e46", size = 52425, upload-time = "2025-08-13T19:32:01.413Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/d9/108cd989a4c0954e60b3cdc86fd2826407702b5375f6dfdab2802e5fed98/rtree-1.4.1-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:d672184298527522d4914d8ae53bf76982b86ca420b0acde9298a7a87d81d4a4", size = 468484, upload-time = "2025-08-13T19:31:50.593Z" }, + { url = "https://files.pythonhosted.org/packages/f3/cf/2710b6fd6b07ea0aef317b29f335790ba6adf06a28ac236078ed9bd8a91d/rtree-1.4.1-py3-none-macosx_11_0_arm64.whl", hash = "sha256:a7e48d805e12011c2cf739a29d6a60ae852fb1de9fc84220bbcef67e6e595d7d", size = 436325, upload-time = "2025-08-13T19:31:52.367Z" }, + { url = "https://files.pythonhosted.org/packages/55/e1/4d075268a46e68db3cac51846eb6a3ab96ed481c585c5a1ad411b3c23aad/rtree-1.4.1-py3-none-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:efa8c4496e31e9ad58ff6c7df89abceac7022d906cb64a3e18e4fceae6b77f65", size = 459789, upload-time = "2025-08-13T19:31:53.926Z" }, + { url = "https://files.pythonhosted.org/packages/d1/75/e5d44be90525cd28503e7f836d077ae6663ec0687a13ba7810b4114b3668/rtree-1.4.1-py3-none-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:12de4578f1b3381a93a655846900be4e3d5f4cd5e306b8b00aa77c1121dc7e8c", size = 507644, upload-time = "2025-08-13T19:31:55.164Z" }, + { url = "https://files.pythonhosted.org/packages/fd/85/b8684f769a142163b52859a38a486493b05bafb4f2fb71d4f945de28ebf9/rtree-1.4.1-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:b558edda52eca3e6d1ee629042192c65e6b7f2c150d6d6cd207ce82f85be3967", size = 1454478, upload-time = "2025-08-13T19:31:56.808Z" }, + { url = "https://files.pythonhosted.org/packages/e9/a4/c2292b95246b9165cc43a0c3757e80995d58bc9b43da5cb47ad6e3535213/rtree-1.4.1-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:f155bc8d6bac9dcd383481dee8c130947a4866db1d16cb6dff442329a038a0dc", size = 1555140, upload-time = "2025-08-13T19:31:58.031Z" }, + { url = "https://files.pythonhosted.org/packages/74/25/5282c8270bfcd620d3e73beb35b40ac4ab00f0a898d98ebeb41ef0989ec8/rtree-1.4.1-py3-none-win_amd64.whl", hash = "sha256:efe125f416fd27150197ab8521158662943a40f87acab8028a1aac4ad667a489", size = 389358, upload-time = "2025-08-13T19:31:59.247Z" }, + { url = "https://files.pythonhosted.org/packages/3f/50/0a9e7e7afe7339bd5e36911f0ceb15fed51945836ed803ae5afd661057fd/rtree-1.4.1-py3-none-win_arm64.whl", hash = "sha256:3d46f55729b28138e897ffef32f7ce93ac335cb67f9120125ad3742a220800f0", size = 355253, upload-time = "2025-08-13T19:32:00.296Z" }, +] + +[[package]] +name = "scipy" +version = "1.17.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/35/48/b992b488d6f299dbe3f11a20b24d3dda3d46f1a635ede1c46b5b17a7b163/scipy-1.17.1-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:35c3a56d2ef83efc372eaec584314bd0ef2e2f0d2adb21c55e6ad5b344c0dcb8", size = 31610954, upload-time = "2026-02-23T00:17:49.855Z" }, + { url = "https://files.pythonhosted.org/packages/b2/02/cf107b01494c19dc100f1d0b7ac3cc08666e96ba2d64db7626066cee895e/scipy-1.17.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:fcb310ddb270a06114bb64bbe53c94926b943f5b7f0842194d585c65eb4edd76", size = 28172662, upload-time = "2026-02-23T00:18:01.64Z" }, + { url = "https://files.pythonhosted.org/packages/cf/a9/599c28631bad314d219cf9ffd40e985b24d603fc8a2f4ccc5ae8419a535b/scipy-1.17.1-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:cc90d2e9c7e5c7f1a482c9875007c095c3194b1cfedca3c2f3291cdc2bc7c086", size = 20344366, upload-time = "2026-02-23T00:18:12.015Z" }, + { url = "https://files.pythonhosted.org/packages/35/f5/906eda513271c8deb5af284e5ef0206d17a96239af79f9fa0aebfe0e36b4/scipy-1.17.1-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:c80be5ede8f3f8eded4eff73cc99a25c388ce98e555b17d31da05287015ffa5b", size = 22704017, upload-time = "2026-02-23T00:18:21.502Z" }, + { url = "https://files.pythonhosted.org/packages/da/34/16f10e3042d2f1d6b66e0428308ab52224b6a23049cb2f5c1756f713815f/scipy-1.17.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e19ebea31758fac5893a2ac360fedd00116cbb7628e650842a6691ba7ca28a21", size = 32927842, upload-time = "2026-02-23T00:18:35.367Z" }, + { url = "https://files.pythonhosted.org/packages/01/8e/1e35281b8ab6d5d72ebe9911edcdffa3f36b04ed9d51dec6dd140396e220/scipy-1.17.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:02ae3b274fde71c5e92ac4d54bc06c42d80e399fec704383dcd99b301df37458", size = 35235890, upload-time = "2026-02-23T00:18:49.188Z" }, + { url = "https://files.pythonhosted.org/packages/c5/5c/9d7f4c88bea6e0d5a4f1bc0506a53a00e9fcb198de372bfe4d3652cef482/scipy-1.17.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8a604bae87c6195d8b1045eddece0514d041604b14f2727bbc2b3020172045eb", size = 35003557, upload-time = "2026-02-23T00:18:54.74Z" }, + { url = "https://files.pythonhosted.org/packages/65/94/7698add8f276dbab7a9de9fb6b0e02fc13ee61d51c7c3f85ac28b65e1239/scipy-1.17.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f590cd684941912d10becc07325a3eeb77886fe981415660d9265c4c418d0bea", size = 37625856, upload-time = "2026-02-23T00:19:00.307Z" }, + { url = "https://files.pythonhosted.org/packages/a2/84/dc08d77fbf3d87d3ee27f6a0c6dcce1de5829a64f2eae85a0ecc1f0daa73/scipy-1.17.1-cp312-cp312-win_amd64.whl", hash = "sha256:41b71f4a3a4cab9d366cd9065b288efc4d4f3c0b37a91a8e0947fb5bd7f31d87", size = 36549682, upload-time = "2026-02-23T00:19:07.67Z" }, + { url = "https://files.pythonhosted.org/packages/bc/98/fe9ae9ffb3b54b62559f52dedaebe204b408db8109a8c66fdd04869e6424/scipy-1.17.1-cp312-cp312-win_arm64.whl", hash = "sha256:f4115102802df98b2b0db3cce5cb9b92572633a1197c77b7553e5203f284a5b3", size = 24547340, upload-time = "2026-02-23T00:19:12.024Z" }, + { url = "https://files.pythonhosted.org/packages/76/27/07ee1b57b65e92645f219b37148a7e7928b82e2b5dbeccecb4dff7c64f0b/scipy-1.17.1-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:5e3c5c011904115f88a39308379c17f91546f77c1667cea98739fe0fccea804c", size = 31590199, upload-time = "2026-02-23T00:19:17.192Z" }, + { url = "https://files.pythonhosted.org/packages/ec/ae/db19f8ab842e9b724bf5dbb7db29302a91f1e55bc4d04b1025d6d605a2c5/scipy-1.17.1-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:6fac755ca3d2c3edcb22f479fceaa241704111414831ddd3bc6056e18516892f", size = 28154001, upload-time = "2026-02-23T00:19:22.241Z" }, + { url = "https://files.pythonhosted.org/packages/5b/58/3ce96251560107b381cbd6e8413c483bbb1228a6b919fa8652b0d4090e7f/scipy-1.17.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:7ff200bf9d24f2e4d5dc6ee8c3ac64d739d3a89e2326ba68aaf6c4a2b838fd7d", size = 20325719, upload-time = "2026-02-23T00:19:26.329Z" }, + { url = "https://files.pythonhosted.org/packages/b2/83/15087d945e0e4d48ce2377498abf5ad171ae013232ae31d06f336e64c999/scipy-1.17.1-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:4b400bdc6f79fa02a4d86640310dde87a21fba0c979efff5248908c6f15fad1b", size = 22683595, upload-time = "2026-02-23T00:19:30.304Z" }, + { url = "https://files.pythonhosted.org/packages/b4/e0/e58fbde4a1a594c8be8114eb4aac1a55bcd6587047efc18a61eb1f5c0d30/scipy-1.17.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2b64ca7d4aee0102a97f3ba22124052b4bd2152522355073580bf4845e2550b6", size = 32896429, upload-time = "2026-02-23T00:19:35.536Z" }, + { url = "https://files.pythonhosted.org/packages/f5/5f/f17563f28ff03c7b6799c50d01d5d856a1d55f2676f537ca8d28c7f627cd/scipy-1.17.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:581b2264fc0aa555f3f435a5944da7504ea3a065d7029ad60e7c3d1ae09c5464", size = 35203952, upload-time = "2026-02-23T00:19:42.259Z" }, + { url = "https://files.pythonhosted.org/packages/8d/a5/9afd17de24f657fdfe4df9a3f1ea049b39aef7c06000c13db1530d81ccca/scipy-1.17.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:beeda3d4ae615106d7094f7e7cef6218392e4465cc95d25f900bebabfded0950", size = 34979063, upload-time = "2026-02-23T00:19:47.547Z" }, + { url = "https://files.pythonhosted.org/packages/8b/13/88b1d2384b424bf7c924f2038c1c409f8d88bb2a8d49d097861dd64a57b2/scipy-1.17.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6609bc224e9568f65064cfa72edc0f24ee6655b47575954ec6339534b2798369", size = 37598449, upload-time = "2026-02-23T00:19:53.238Z" }, + { url = "https://files.pythonhosted.org/packages/35/e5/d6d0e51fc888f692a35134336866341c08655d92614f492c6860dc45bb2c/scipy-1.17.1-cp313-cp313-win_amd64.whl", hash = "sha256:37425bc9175607b0268f493d79a292c39f9d001a357bebb6b88fdfaff13f6448", size = 36510943, upload-time = "2026-02-23T00:20:50.89Z" }, + { url = "https://files.pythonhosted.org/packages/2a/fd/3be73c564e2a01e690e19cc618811540ba5354c67c8680dce3281123fb79/scipy-1.17.1-cp313-cp313-win_arm64.whl", hash = "sha256:5cf36e801231b6a2059bf354720274b7558746f3b1a4efb43fcf557ccd484a87", size = 24545621, upload-time = "2026-02-23T00:20:55.871Z" }, + { url = "https://files.pythonhosted.org/packages/6f/6b/17787db8b8114933a66f9dcc479a8272e4b4da75fe03b0c282f7b0ade8cd/scipy-1.17.1-cp313-cp313t-macosx_10_14_x86_64.whl", hash = "sha256:d59c30000a16d8edc7e64152e30220bfbd724c9bbb08368c054e24c651314f0a", size = 31936708, upload-time = "2026-02-23T00:19:58.694Z" }, + { url = "https://files.pythonhosted.org/packages/38/2e/524405c2b6392765ab1e2b722a41d5da33dc5c7b7278184a8ad29b6cb206/scipy-1.17.1-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:010f4333c96c9bb1a4516269e33cb5917b08ef2166d5556ca2fd9f082a9e6ea0", size = 28570135, upload-time = "2026-02-23T00:20:03.934Z" }, + { url = "https://files.pythonhosted.org/packages/fd/c3/5bd7199f4ea8556c0c8e39f04ccb014ac37d1468e6cfa6a95c6b3562b76e/scipy-1.17.1-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:2ceb2d3e01c5f1d83c4189737a42d9cb2fc38a6eeed225e7515eef71ad301dce", size = 20741977, upload-time = "2026-02-23T00:20:07.935Z" }, + { url = "https://files.pythonhosted.org/packages/d9/b8/8ccd9b766ad14c78386599708eb745f6b44f08400a5fd0ade7cf89b6fc93/scipy-1.17.1-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:844e165636711ef41f80b4103ed234181646b98a53c8f05da12ca5ca289134f6", size = 23029601, upload-time = "2026-02-23T00:20:12.161Z" }, + { url = "https://files.pythonhosted.org/packages/6d/a0/3cb6f4d2fb3e17428ad2880333cac878909ad1a89f678527b5328b93c1d4/scipy-1.17.1-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:158dd96d2207e21c966063e1635b1063cd7787b627b6f07305315dd73d9c679e", size = 33019667, upload-time = "2026-02-23T00:20:17.208Z" }, + { url = "https://files.pythonhosted.org/packages/f3/c3/2d834a5ac7bf3a0c806ad1508efc02dda3c8c61472a56132d7894c312dea/scipy-1.17.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74cbb80d93260fe2ffa334efa24cb8f2f0f622a9b9febf8b483c0b865bfb3475", size = 35264159, upload-time = "2026-02-23T00:20:23.087Z" }, + { url = "https://files.pythonhosted.org/packages/4d/77/d3ed4becfdbd217c52062fafe35a72388d1bd82c2d0ba5ca19d6fcc93e11/scipy-1.17.1-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:dbc12c9f3d185f5c737d801da555fb74b3dcfa1a50b66a1a93e09190f41fab50", size = 35102771, upload-time = "2026-02-23T00:20:28.636Z" }, + { url = "https://files.pythonhosted.org/packages/bd/12/d19da97efde68ca1ee5538bb261d5d2c062f0c055575128f11a2730e3ac1/scipy-1.17.1-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:94055a11dfebe37c656e70317e1996dc197e1a15bbcc351bcdd4610e128fe1ca", size = 37665910, upload-time = "2026-02-23T00:20:34.743Z" }, + { url = "https://files.pythonhosted.org/packages/06/1c/1172a88d507a4baaf72c5a09bb6c018fe2ae0ab622e5830b703a46cc9e44/scipy-1.17.1-cp313-cp313t-win_amd64.whl", hash = "sha256:e30bdeaa5deed6bc27b4cc490823cd0347d7dae09119b8803ae576ea0ce52e4c", size = 36562980, upload-time = "2026-02-23T00:20:40.575Z" }, + { url = "https://files.pythonhosted.org/packages/70/b0/eb757336e5a76dfa7911f63252e3b7d1de00935d7705cf772db5b45ec238/scipy-1.17.1-cp313-cp313t-win_arm64.whl", hash = "sha256:a720477885a9d2411f94a93d16f9d89bad0f28ca23c3f8daa521e2dcc3f44d49", size = 24856543, upload-time = "2026-02-23T00:20:45.313Z" }, +] + +[[package]] +name = "sentry-sdk" +version = "2.58.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/26/b3/fb8291170d0e844173164709fc0fa0c221ed75a5da740c8746f2a83b4eb1/sentry_sdk-2.58.0.tar.gz", hash = "sha256:c1144d947352d54e5b7daa63596d9f848adf684989c06c4f5a659f0c85a18f6f", size = 438764, upload-time = "2026-04-13T17:23:26.265Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fa/eb/d875669993b762556ae8b2efd86219943b4c0864d22204d622a9aee3052b/sentry_sdk-2.58.0-py2.py3-none-any.whl", hash = "sha256:688d1c704ddecf382ea3326f21a67453d4caa95592d722b7c780a36a9d23109e", size = 460919, upload-time = "2026-04-13T17:23:24.675Z" }, +] + +[[package]] +name = "setuptools" +version = "81.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0d/1c/73e719955c59b8e424d015ab450f51c0af856ae46ea2da83eba51cc88de1/setuptools-81.0.0.tar.gz", hash = "sha256:487b53915f52501f0a79ccfd0c02c165ffe06631443a886740b91af4b7a5845a", size = 1198299, upload-time = "2026-02-06T21:10:39.601Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e1/e3/c164c88b2e5ce7b24d667b9bd83589cf4f3520d97cad01534cd3c4f55fdb/setuptools-81.0.0-py3-none-any.whl", hash = "sha256:fdd925d5c5d9f62e4b74b30d6dd7828ce236fd6ed998a08d81de62ce5a6310d6", size = 1062021, upload-time = "2026-02-06T21:10:37.175Z" }, +] + +[[package]] +name = "shapely" +version = "2.1.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4d/bc/0989043118a27cccb4e906a46b7565ce36ca7b57f5a18b78f4f1b0f72d9d/shapely-2.1.2.tar.gz", hash = "sha256:2ed4ecb28320a433db18a5bf029986aa8afcfd740745e78847e330d5d94922a9", size = 315489, upload-time = "2025-09-24T13:51:41.432Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/c0/f3b6453cf2dfa99adc0ba6675f9aaff9e526d2224cbd7ff9c1a879238693/shapely-2.1.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:fe2533caae6a91a543dec62e8360fe86ffcdc42a7c55f9dfd0128a977a896b94", size = 1833550, upload-time = "2025-09-24T13:50:30.019Z" }, + { url = "https://files.pythonhosted.org/packages/86/07/59dee0bc4b913b7ab59ab1086225baca5b8f19865e6101db9ebb7243e132/shapely-2.1.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ba4d1333cc0bc94381d6d4308d2e4e008e0bd128bdcff5573199742ee3634359", size = 1643556, upload-time = "2025-09-24T13:50:32.291Z" }, + { url = "https://files.pythonhosted.org/packages/26/29/a5397e75b435b9895cd53e165083faed5d12fd9626eadec15a83a2411f0f/shapely-2.1.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0bd308103340030feef6c111d3eb98d50dc13feea33affc8a6f9fa549e9458a3", size = 2988308, upload-time = "2025-09-24T13:50:33.862Z" }, + { url = "https://files.pythonhosted.org/packages/b9/37/e781683abac55dde9771e086b790e554811a71ed0b2b8a1e789b7430dd44/shapely-2.1.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1e7d4d7ad262a48bb44277ca12c7c78cb1b0f56b32c10734ec9a1d30c0b0c54b", size = 3099844, upload-time = "2025-09-24T13:50:35.459Z" }, + { url = "https://files.pythonhosted.org/packages/d8/f3/9876b64d4a5a321b9dc482c92bb6f061f2fa42131cba643c699f39317cb9/shapely-2.1.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e9eddfe513096a71896441a7c37db72da0687b34752c4e193577a145c71736fc", size = 3988842, upload-time = "2025-09-24T13:50:37.478Z" }, + { url = "https://files.pythonhosted.org/packages/d1/a0/704c7292f7014c7e74ec84eddb7b109e1fbae74a16deae9c1504b1d15565/shapely-2.1.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:980c777c612514c0cf99bc8a9de6d286f5e186dcaf9091252fcd444e5638193d", size = 4152714, upload-time = "2025-09-24T13:50:39.9Z" }, + { url = "https://files.pythonhosted.org/packages/53/46/319c9dc788884ad0785242543cdffac0e6530e4d0deb6c4862bc4143dcf3/shapely-2.1.2-cp312-cp312-win32.whl", hash = "sha256:9111274b88e4d7b54a95218e243282709b330ef52b7b86bc6aaf4f805306f454", size = 1542745, upload-time = "2025-09-24T13:50:41.414Z" }, + { url = "https://files.pythonhosted.org/packages/ec/bf/cb6c1c505cb31e818e900b9312d514f381fbfa5c4363edfce0fcc4f8c1a4/shapely-2.1.2-cp312-cp312-win_amd64.whl", hash = "sha256:743044b4cfb34f9a67205cee9279feaf60ba7d02e69febc2afc609047cb49179", size = 1722861, upload-time = "2025-09-24T13:50:43.35Z" }, + { url = "https://files.pythonhosted.org/packages/c3/90/98ef257c23c46425dc4d1d31005ad7c8d649fe423a38b917db02c30f1f5a/shapely-2.1.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:b510dda1a3672d6879beb319bc7c5fd302c6c354584690973c838f46ec3e0fa8", size = 1832644, upload-time = "2025-09-24T13:50:44.886Z" }, + { url = "https://files.pythonhosted.org/packages/6d/ab/0bee5a830d209adcd3a01f2d4b70e587cdd9fd7380d5198c064091005af8/shapely-2.1.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8cff473e81017594d20ec55d86b54bc635544897e13a7cfc12e36909c5309a2a", size = 1642887, upload-time = "2025-09-24T13:50:46.735Z" }, + { url = "https://files.pythonhosted.org/packages/2d/5e/7d7f54ba960c13302584c73704d8c4d15404a51024631adb60b126a4ae88/shapely-2.1.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fe7b77dc63d707c09726b7908f575fc04ff1d1ad0f3fb92aec212396bc6cfe5e", size = 2970931, upload-time = "2025-09-24T13:50:48.374Z" }, + { url = "https://files.pythonhosted.org/packages/f2/a2/83fc37e2a58090e3d2ff79175a95493c664bcd0b653dd75cb9134645a4e5/shapely-2.1.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7ed1a5bbfb386ee8332713bf7508bc24e32d24b74fc9a7b9f8529a55db9f4ee6", size = 3082855, upload-time = "2025-09-24T13:50:50.037Z" }, + { url = "https://files.pythonhosted.org/packages/44/2b/578faf235a5b09f16b5f02833c53822294d7f21b242f8e2d0cf03fb64321/shapely-2.1.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a84e0582858d841d54355246ddfcbd1fce3179f185da7470f41ce39d001ee1af", size = 3979960, upload-time = "2025-09-24T13:50:51.74Z" }, + { url = "https://files.pythonhosted.org/packages/4d/04/167f096386120f692cc4ca02f75a17b961858997a95e67a3cb6a7bbd6b53/shapely-2.1.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:dc3487447a43d42adcdf52d7ac73804f2312cbfa5d433a7d2c506dcab0033dfd", size = 4142851, upload-time = "2025-09-24T13:50:53.49Z" }, + { url = "https://files.pythonhosted.org/packages/48/74/fb402c5a6235d1c65a97348b48cdedb75fb19eca2b1d66d04969fc1c6091/shapely-2.1.2-cp313-cp313-win32.whl", hash = "sha256:9c3a3c648aedc9f99c09263b39f2d8252f199cb3ac154fadc173283d7d111350", size = 1541890, upload-time = "2025-09-24T13:50:55.337Z" }, + { url = "https://files.pythonhosted.org/packages/41/47/3647fe7ad990af60ad98b889657a976042c9988c2807cf322a9d6685f462/shapely-2.1.2-cp313-cp313-win_amd64.whl", hash = "sha256:ca2591bff6645c216695bdf1614fca9c82ea1144d4a7591a466fef64f28f0715", size = 1722151, upload-time = "2025-09-24T13:50:57.153Z" }, + { url = "https://files.pythonhosted.org/packages/3c/49/63953754faa51ffe7d8189bfbe9ca34def29f8c0e34c67cbe2a2795f269d/shapely-2.1.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:2d93d23bdd2ed9dc157b46bc2f19b7da143ca8714464249bef6771c679d5ff40", size = 1834130, upload-time = "2025-09-24T13:50:58.49Z" }, + { url = "https://files.pythonhosted.org/packages/7f/ee/dce001c1984052970ff60eb4727164892fb2d08052c575042a47f5a9e88f/shapely-2.1.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:01d0d304b25634d60bd7cf291828119ab55a3bab87dc4af1e44b07fb225f188b", size = 1642802, upload-time = "2025-09-24T13:50:59.871Z" }, + { url = "https://files.pythonhosted.org/packages/da/e7/fc4e9a19929522877fa602f705706b96e78376afb7fad09cad5b9af1553c/shapely-2.1.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8d8382dd120d64b03698b7298b89611a6ea6f55ada9d39942838b79c9bc89801", size = 3018460, upload-time = "2025-09-24T13:51:02.08Z" }, + { url = "https://files.pythonhosted.org/packages/a1/18/7519a25db21847b525696883ddc8e6a0ecaa36159ea88e0fef11466384d0/shapely-2.1.2-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:19efa3611eef966e776183e338b2d7ea43569ae99ab34f8d17c2c054d3205cc0", size = 3095223, upload-time = "2025-09-24T13:51:04.472Z" }, + { url = "https://files.pythonhosted.org/packages/48/de/b59a620b1f3a129c3fecc2737104a0a7e04e79335bd3b0a1f1609744cf17/shapely-2.1.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:346ec0c1a0fcd32f57f00e4134d1200e14bf3f5ae12af87ba83ca275c502498c", size = 4030760, upload-time = "2025-09-24T13:51:06.455Z" }, + { url = "https://files.pythonhosted.org/packages/96/b3/c6655ee7232b417562bae192ae0d3ceaadb1cc0ffc2088a2ddf415456cc2/shapely-2.1.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:6305993a35989391bd3476ee538a5c9a845861462327efe00dd11a5c8c709a99", size = 4170078, upload-time = "2025-09-24T13:51:08.584Z" }, + { url = "https://files.pythonhosted.org/packages/a0/8e/605c76808d73503c9333af8f6cbe7e1354d2d238bda5f88eea36bfe0f42a/shapely-2.1.2-cp313-cp313t-win32.whl", hash = "sha256:c8876673449f3401f278c86eb33224c5764582f72b653a415d0e6672fde887bf", size = 1559178, upload-time = "2025-09-24T13:51:10.73Z" }, + { url = "https://files.pythonhosted.org/packages/36/f7/d317eb232352a1f1444d11002d477e54514a4a6045536d49d0c59783c0da/shapely-2.1.2-cp313-cp313t-win_amd64.whl", hash = "sha256:4a44bc62a10d84c11a7a3d7c1c4fe857f7477c3506e24c9062da0db0ae0c449c", size = 1739756, upload-time = "2025-09-24T13:51:12.105Z" }, +] + +[[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 = "smmap" +version = "5.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1f/ea/49c993d6dfdd7338c9b1000a0f36817ed7ec84577ae2e52f890d1a4ff909/smmap-5.0.3.tar.gz", hash = "sha256:4d9debb8b99007ae47165abc08670bd74cb74b5227dda7f643eccc4e9eb5642c", size = 22506, upload-time = "2026-03-09T03:43:26.1Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/d4/59e74daffcb57a07668852eeeb6035af9f32cbfd7a1d2511f17d2fe6a738/smmap-5.0.3-py3-none-any.whl", hash = "sha256:c106e05d5a61449cf6ba9a1e650227ecfb141590d2a98412103ff35d89fc7b2f", size = 24390, upload-time = "2026-03-09T03:43:24.361Z" }, +] + +[[package]] +name = "svg-path" +version = "7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/66/b9/649abbe870842c185b12920e937e9b95d4c2b18de50af98d2c140df3e179/svg_path-7.0.tar.gz", hash = "sha256:9037486957cb1dcf4375ef42206499a47c111b8ffcbac6e3e55f9d079d875bb0", size = 23552, upload-time = "2025-07-06T15:20:40.823Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3a/83/4f5b250220e1a5acd31345a5ec1c95a7769725d0d8135276f399f44062f8/svg_path-7.0-py2.py3-none-any.whl", hash = "sha256:447cb1e16a95acea2dd867fe737fa99cb75d587b4fc64dbee709a8dd6891ad9c", size = 18208, upload-time = "2025-07-06T15:20:39.59Z" }, +] + +[[package]] +name = "sympy" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mpmath" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, +] + +[[package]] +name = "tensorboard" +version = "2.20.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "grpcio" }, + { name = "markdown" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pillow" }, + { name = "protobuf" }, + { name = "setuptools" }, + { name = "tensorboard-data-server" }, + { name = "werkzeug" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/9c/d9/a5db55f88f258ac669a92858b70a714bbbd5acd993820b41ec4a96a4d77f/tensorboard-2.20.0-py3-none-any.whl", hash = "sha256:9dc9f978cb84c0723acf9a345d96c184f0293d18f166bb8d59ee098e6cfaaba6", size = 5525680, upload-time = "2025-07-17T19:20:49.638Z" }, +] + +[[package]] +name = "tensorboard-data-server" +version = "0.7.2" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7a/13/e503968fefabd4c6b2650af21e110aa8466fe21432cd7c43a84577a89438/tensorboard_data_server-0.7.2-py3-none-any.whl", hash = "sha256:7e0610d205889588983836ec05dc098e80f97b7e7bbff7e994ebb78f578d0ddb", size = 2356, upload-time = "2023-10-23T21:23:32.16Z" }, + { url = "https://files.pythonhosted.org/packages/b7/85/dabeaf902892922777492e1d253bb7e1264cadce3cea932f7ff599e53fea/tensorboard_data_server-0.7.2-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:9fe5d24221b29625dbc7328b0436ca7fc1c23de4acf4d272f1180856e32f9f60", size = 4823598, upload-time = "2023-10-23T21:23:33.714Z" }, + { url = "https://files.pythonhosted.org/packages/73/c6/825dab04195756cf8ff2e12698f22513b3db2f64925bdd41671bfb33aaa5/tensorboard_data_server-0.7.2-py3-none-manylinux_2_31_x86_64.whl", hash = "sha256:ef687163c24185ae9754ed5650eb5bc4d84ff257aabdc33f0cc6f74d8ba54530", size = 6590363, upload-time = "2023-10-23T21:23:35.583Z" }, +] + +[[package]] +name = "tensordict" +version = "0.12.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cloudpickle" }, + { name = "importlib-metadata" }, + { name = "numpy" }, + { name = "orjson", marker = "python_full_version < '3.13'" }, + { name = "packaging" }, + { name = "pyvers" }, + { name = "torch" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/31/fd/034d044b4019873ed64c55edc0932b39ae0c7d2e55177a2bd8da1dc1c1f2/tensordict-0.12.2-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:9fc9a5029840d0cbbfcc2d6b623577d6a63a6322b13eb263c7c0c5c13d907e56", size = 889351, upload-time = "2026-04-20T15:11:28.702Z" }, + { url = "https://files.pythonhosted.org/packages/dc/88/26a3af5dc0a9ade8d10d32c174ba74d6e1bf9641b71f49cce311e8426407/tensordict-0.12.2-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:1642121b3ae54106e246c58907e6a60a1b32ab36679167c0fcef336760c19ff2", size = 532720, upload-time = "2026-04-20T15:11:30.326Z" }, + { url = "https://files.pythonhosted.org/packages/94/54/f33d016855d076387141a96b805e8e4d394139c94f1a5c380e7c187acc62/tensordict-0.12.2-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:1815a93ae74f9c8d2d530a8d5de920b2aef9aa07da7f49200e74aea3c6a49894", size = 536771, upload-time = "2026-04-20T15:11:32.34Z" }, + { url = "https://files.pythonhosted.org/packages/17/2a/418bab656a7af277cf1bdd725219a881d842178ff31f1756c298252a4bfe/tensordict-0.12.2-cp312-cp312-win_amd64.whl", hash = "sha256:69b2c4a07f5226076753b9bc6d45355376d612da8817bde1f063a9e9c7a9a28f", size = 586030, upload-time = "2026-04-20T15:11:33.909Z" }, + { url = "https://files.pythonhosted.org/packages/4d/00/bd86f3df83d4718a6d57768cffbe235440f52cb7caafa77d19c3661ec5a2/tensordict-0.12.2-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:ce53dd911d63719edd5462e1d6dfae4bd55e4b5fa5bceb7fac9b8b0749a715a5", size = 889359, upload-time = "2026-04-20T15:11:35.593Z" }, + { url = "https://files.pythonhosted.org/packages/ef/61/4b51ab1892155fa6fc3373773cdea7beb56e5636a6484459dd7452636bca/tensordict-0.12.2-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:e005a04d00b499a1a36883338145ae014ddd53a9498e369535d4c499c8867928", size = 532982, upload-time = "2026-04-20T15:11:37.25Z" }, + { url = "https://files.pythonhosted.org/packages/56/49/a851c2c610ed6d08714d4c6af91287cfb250a70fa166678d09f48e532cea/tensordict-0.12.2-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:49b575a39dc1a8de138e6e519329b55eae39fba721ff43aa4e0c08afcacd5fe3", size = 536753, upload-time = "2026-04-20T15:11:38.707Z" }, + { url = "https://files.pythonhosted.org/packages/14/31/14da5697d6e57740a507fdb0c2daa424f67603647071e123b9a1f5293f00/tensordict-0.12.2-cp313-cp313-win_amd64.whl", hash = "sha256:2710b7ce7730c544d2519b0b466a0d47a61319e552c49da54d454d41ccef452f", size = 586005, upload-time = "2026-04-20T15:11:40.365Z" }, + { url = "https://files.pythonhosted.org/packages/2a/2e/b9509652ddd69de4b738cef8f246072667fc51a91be026f005f3e666657d/tensordict-0.12.2-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:70b185f0f9545f5e79d64383498a933b780cd14d017b447556e4d4ed1e0f3e33", size = 894783, upload-time = "2026-04-20T15:11:42.12Z" }, + { url = "https://files.pythonhosted.org/packages/d3/d3/41a21801bbc1c6cf6374c4f7271904815095a5b3375f22c14d0f7e02050e/tensordict-0.12.2-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:0c881da6d48189357ab414f9cb3394a6d0513076b2287c3e7f9a47e5d0ab1730", size = 534421, upload-time = "2026-04-20T15:11:43.496Z" }, + { url = "https://files.pythonhosted.org/packages/d0/d3/828793ad818935b300fb61eb0c9041c572bb6f8d124cef43e6323a6f6b4d/tensordict-0.12.2-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:8294507ea68b37c342087113f651bd36f823b805bd7cabe9440c587d507fc744", size = 538294, upload-time = "2026-04-20T15:11:44.814Z" }, + { url = "https://files.pythonhosted.org/packages/d8/eb/43e87ba618ed1844e5a537258381966e12fc0b032bfb57d617cb7395d818/tensordict-0.12.2-cp313-cp313t-win_amd64.whl", hash = "sha256:3e1a93bffe9d459616724327c8f3e0b05d63737db94232d69913ffa5af2b81d1", size = 596851, upload-time = "2026-04-20T15:11:46.292Z" }, +] + +[[package]] +name = "torch" +version = "2.11.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-bindings", marker = "sys_platform == 'linux'" }, + { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" }, + { name = "setuptools" }, + { name = "sympy" }, + { name = "triton", marker = "sys_platform == 'linux'" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/6f/8b/69e3008d78e5cee2b30183340cc425081b78afc5eff3d080daab0adda9aa/torch-2.11.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4b5866312ee6e52ea625cd211dcb97d6a2cdc1131a5f15cc0d87eec948f6dd34", size = 80606338, upload-time = "2026-03-23T18:11:34.781Z" }, + { url = "https://files.pythonhosted.org/packages/13/16/42e5915ebe4868caa6bac83a8ed59db57f12e9a61b7d749d584776ed53d5/torch-2.11.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:f99924682ef0aa6a4ab3b1b76f40dc6e273fca09f367d15a524266db100a723f", size = 419731115, upload-time = "2026-03-23T18:11:06.944Z" }, + { url = "https://files.pythonhosted.org/packages/1a/c9/82638ef24d7877510f83baf821f5619a61b45568ce21c0a87a91576510aa/torch-2.11.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:0f68f4ac6d95d12e896c3b7a912b5871619542ec54d3649cf48cc1edd4dd2756", size = 530712279, upload-time = "2026-03-23T18:10:31.481Z" }, + { url = "https://files.pythonhosted.org/packages/1c/ff/6756f1c7ee302f6d202120e0f4f05b432b839908f9071157302cedfc5232/torch-2.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:fbf39280699d1b869f55eac536deceaa1b60bd6788ba74f399cc67e60a5fab10", size = 114556047, upload-time = "2026-03-23T18:10:55.931Z" }, + { url = "https://files.pythonhosted.org/packages/87/89/5ea6722763acee56b045435fb84258db7375c48165ec8be7880ab2b281c5/torch-2.11.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1e6debd97ccd3205bbb37eb806a9d8219e1139d15419982c09e23ef7d4369d18", size = 80606801, upload-time = "2026-03-23T18:10:18.649Z" }, + { url = "https://files.pythonhosted.org/packages/32/d1/8ed2173589cbfe744ed54e5a73efc107c0085ba5777ee93a5f4c1ab90553/torch-2.11.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:63a68fa59de8f87acc7e85a5478bb2dddbb3392b7593ec3e78827c793c4b73fd", size = 419732382, upload-time = "2026-03-23T18:08:30.835Z" }, + { url = "https://files.pythonhosted.org/packages/3d/e1/b73f7c575a4b8f87a5928f50a1e35416b5e27295d8be9397d5293e7e8d4c/torch-2.11.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:cc89b9b173d9adfab59fd227f0ab5e5516d9a52b658ae41d64e59d2e55a418db", size = 530711509, upload-time = "2026-03-23T18:08:47.213Z" }, + { url = "https://files.pythonhosted.org/packages/66/82/3e3fcdd388fbe54e29fd3f991f36846ff4ac90b0d0181e9c8f7236565f82/torch-2.11.0-cp313-cp313-win_amd64.whl", hash = "sha256:4dda3b3f52d121063a731ddb835f010dc137b920d7fec2778e52f60d8e4bf0cd", size = 114555842, upload-time = "2026-03-23T18:09:52.111Z" }, + { url = "https://files.pythonhosted.org/packages/db/38/8ac78069621b8c2b4979c2f96dc8409ef5e9c4189f6aac629189a78677ca/torch-2.11.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:8b394322f49af4362d4f80e424bcaca7efcd049619af03a4cf4501520bdf0fb4", size = 80959574, upload-time = "2026-03-23T18:10:14.214Z" }, + { url = "https://files.pythonhosted.org/packages/6d/6c/56bfb37073e7136e6dd86bfc6af7339946dd684e0ecf2155ac0eee687ae1/torch-2.11.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:2658f34ce7e2dabf4ec73b45e2ca68aedad7a5be87ea756ad656eaf32bf1e1ea", size = 419732324, upload-time = "2026-03-23T18:09:36.604Z" }, + { url = "https://files.pythonhosted.org/packages/07/f4/1b666b6d61d3394cca306ea543ed03a64aad0a201b6cd159f1d41010aeb1/torch-2.11.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:98bb213c3084cfe176302949bdc360074b18a9da7ab59ef2edc9d9f742504778", size = 530596026, upload-time = "2026-03-23T18:09:20.842Z" }, + { url = "https://files.pythonhosted.org/packages/48/6b/30d1459fa7e4b67e9e3fe1685ca1d8bb4ce7c62ef436c3a615963c6c866c/torch-2.11.0-cp313-cp313t-win_amd64.whl", hash = "sha256:a97b94bbf62992949b4730c6cd2cc9aee7b335921ee8dc207d930f2ed09ae2db", size = 114793702, upload-time = "2026-03-23T18:09:47.304Z" }, +] + +[[package]] +name = "torchrunx" +version = "0.3.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cloudpickle" }, + { name = "fabric" }, + { name = "numpy" }, + { name = "torch" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/65/60/ce6ccaf618e56775905e75e3fe7f2c8adfb61916d2946e854850c1d19a0d/torchrunx-0.3.4.tar.gz", hash = "sha256:6f2333fa17f7ef1f43f6c65d2b008b8479b29d972a8ed209da613d830dffdc45", size = 41312, upload-time = "2025-11-19T02:36:13.54Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2c/0d/9f5e24043f2562fd4dc15d04614cb5f3b4a1ffe327cb040b9f95b03fa84d/torchrunx-0.3.4-py3-none-any.whl", hash = "sha256:a157ec139f5a0bdfaa5ece50d987ff0a3212a4657d791367f11fdd383ccdbd5b", size = 34172, upload-time = "2025-11-19T02:36:12.354Z" }, +] + +[[package]] +name = "torchvision" +version = "0.26.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "pillow" }, + { name = "torch" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/ae/e7/56b47cc3b132aea90ccce22bcb8975dec688b002150012acc842846039d0/torchvision-0.26.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c409e1c3fdebec7a3834465086dbda8bf7680eff79abf7fd2f10c6b59520a7a4", size = 1863502, upload-time = "2026-03-23T18:12:57.326Z" }, + { url = "https://files.pythonhosted.org/packages/f4/ec/5c31c92c08b65662fe9604a4067ae8232582805949f11ddc042cebe818ed/torchvision-0.26.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:406557718e62fdf10f5706e88d8a5ec000f872da913bf629aab9297622585547", size = 7767944, upload-time = "2026-03-23T18:12:42.805Z" }, + { url = "https://files.pythonhosted.org/packages/f5/d8/cb6ccda1a1f35a6597645818641701207b3e8e13553e75fce5d86bac74b2/torchvision-0.26.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:d61a5abb6b42a0c0c311996c2ac4b83a94418a97182c83b055a2a4ae985e05aa", size = 7522205, upload-time = "2026-03-23T18:12:54.654Z" }, + { url = "https://files.pythonhosted.org/packages/1c/a9/c272623a0f735c35f0f6cd6dc74784d4f970e800cf063bb76687895a2ab9/torchvision-0.26.0-cp312-cp312-win_amd64.whl", hash = "sha256:7993c01648e7c61d191b018e84d38fe0825c8fcb2720cd0f37caf7ba14404aa1", size = 4255155, upload-time = "2026-03-23T18:12:32.652Z" }, + { url = "https://files.pythonhosted.org/packages/da/80/0762f77f53605d10c9477be39bb47722cc8e383bbbc2531471ce0e396c07/torchvision-0.26.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:5d63dd43162691258b1b3529b9041bac7d54caa37eae0925f997108268cbf7c4", size = 1860809, upload-time = "2026-03-23T18:12:47.629Z" }, + { url = "https://files.pythonhosted.org/packages/e6/81/0b3e58d1478c660a5af4268713486b2df7203f35abd9195fea87348a5178/torchvision-0.26.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:a39c7a26538c41fda453f9a9692b5ff9b35a5437db1d94f3027f6f509c160eac", size = 7727494, upload-time = "2026-03-23T18:12:46.062Z" }, + { url = "https://files.pythonhosted.org/packages/b6/dc/d9ab5d29115aa05e12e30f1397a3eeae1d88a511241dc3bce48dc4342675/torchvision-0.26.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:b7e6213620bbf97742e5f79832f9e9d769e6cf0f744c5b53dad80b76db633691", size = 7521747, upload-time = "2026-03-23T18:12:36.815Z" }, + { url = "https://files.pythonhosted.org/packages/a9/1b/f1bc86a918c5f6feab1eeff11982e2060f4704332e96185463d27855bdf5/torchvision-0.26.0-cp313-cp313-win_amd64.whl", hash = "sha256:4280c35ec8cba1fcc8294fb87e136924708726864c379e4c54494797d86bc474", size = 4319880, upload-time = "2026-03-23T18:12:38.168Z" }, + { url = "https://files.pythonhosted.org/packages/66/28/b4ad0a723ed95b003454caffcc41894b34bd8379df340848cae2c33871de/torchvision-0.26.0-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:358fc4726d0c08615b6d83b3149854f11efb2a564ed1acb6fce882e151412d23", size = 1951973, upload-time = "2026-03-23T18:12:48.781Z" }, + { url = "https://files.pythonhosted.org/packages/71/e2/7a89096e6cf2f3336353b5338ba925e0addf9d8601920340e6bdf47e8eb3/torchvision-0.26.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:3daf9cc149cf3cdcbd4df9c59dae69ffca86c6823250442c3bbfd63fc2e26c61", size = 7728679, upload-time = "2026-03-23T18:12:26.196Z" }, + { url = "https://files.pythonhosted.org/packages/69/1d/4e1eebc17d18ce080a11dcf3df3f8f717f0efdfa00983f06e8ba79259f61/torchvision-0.26.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:82c3965eca27e86a316e31e4c3e5a16d353e0bcbe0ef8efa2e66502c54493c4b", size = 7609138, upload-time = "2026-03-23T18:12:35.327Z" }, + { url = "https://files.pythonhosted.org/packages/f3/a4/f1155e943ae5b32400d7000adc81c79bb0392b16ceb33bcf13e02e48cced/torchvision-0.26.0-cp313-cp313t-win_amd64.whl", hash = "sha256:ebc043cc5a4f0bf22e7680806dbba37ffb19e70f6953bbb44ed1a90aeb5c9bea", size = 4248202, upload-time = "2026-03-23T18:12:41.423Z" }, +] + +[[package]] +name = "tqdm" +version = "4.67.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/09/a9/6ba95a270c6f1fbcd8dac228323f2777d886cb206987444e4bce66338dd4/tqdm-4.67.3.tar.gz", hash = "sha256:7d825f03f89244ef73f1d4ce193cb1774a8179fd96f31d7e1dcde62092b960bb", size = 169598, upload-time = "2026-02-03T17:35:53.048Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/16/e1/3079a9ff9b8e11b846c6ac5c8b5bfb7ff225eee721825310c91b3b50304f/tqdm-4.67.3-py3-none-any.whl", hash = "sha256:ee1e4c0e59148062281c49d80b25b67771a127c85fc9676d3be5f243206826bf", size = 78374, upload-time = "2026-02-03T17:35:50.982Z" }, +] + +[[package]] +name = "trimesh" +version = "4.11.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0d/bf/53b69f3b6708c20ceb4d1d1250c7dc205733eb646659e5e55771f76ffabd/trimesh-4.11.5.tar.gz", hash = "sha256:b90e6cdd6ada51c52d4a7d32947f4ce44b6751c5b7cab2b04e271ecea1e397d3", size = 836449, upload-time = "2026-03-25T01:08:24.216Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/83/72e812f772daee66651f468c7b2535fa05eac27db26df7e614cae823c832/trimesh-4.11.5-py3-none-any.whl", hash = "sha256:b225a94c8af79569f7167ca7eaaab4fd05c260da58a075599453d655835258ef", size = 740833, upload-time = "2026-03-25T01:08:21.397Z" }, +] + +[package.optional-dependencies] +easy = [ + { name = "charset-normalizer" }, + { name = "colorlog" }, + { name = "embreex", marker = "platform_machine == 'x86_64'" }, + { name = "httpx" }, + { name = "jsonschema" }, + { name = "lxml" }, + { name = "manifold3d" }, + { name = "mapbox-earcut" }, + { name = "networkx" }, + { name = "pillow" }, + { name = "pycollada" }, + { name = "rtree" }, + { name = "scipy" }, + { name = "shapely" }, + { name = "svg-path" }, + { name = "vhacdx" }, + { name = "xxhash" }, +] + +[[package]] +name = "triton" +version = "3.6.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/17/5d/08201db32823bdf77a0e2b9039540080b2e5c23a20706ddba942924ebcd6/triton-3.6.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:374f52c11a711fd062b4bfbb201fd9ac0a5febd28a96fb41b4a0f51dde3157f4", size = 176128243, upload-time = "2026-01-20T16:16:07.857Z" }, + { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, + { url = "https://files.pythonhosted.org/packages/3c/12/34d71b350e89a204c2c7777a9bba0dcf2f19a5bfdd70b57c4dbc5ffd7154/triton-3.6.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:448e02fe6dc898e9e5aa89cf0ee5c371e99df5aa5e8ad976a80b93334f3494fd", size = 176133521, upload-time = "2026-01-20T16:16:13.321Z" }, + { url = "https://files.pythonhosted.org/packages/f9/0b/37d991d8c130ce81a8728ae3c25b6e60935838e9be1b58791f5997b24a54/triton-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10c7f76c6e72d2ef08df639e3d0d30729112f47a56b0c81672edc05ee5116ac9", size = 188289450, upload-time = "2026-01-20T16:00:49.136Z" }, + { url = "https://files.pythonhosted.org/packages/ce/4e/41b0c8033b503fd3cfcd12392cdd256945026a91ff02452bef40ec34bee7/triton-3.6.0-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1722e172d34e32abc3eb7711d0025bb69d7959ebea84e3b7f7a341cd7ed694d6", size = 176276087, upload-time = "2026-01-20T16:16:18.989Z" }, + { url = "https://files.pythonhosted.org/packages/35/f8/9c66bfc55361ec6d0e4040a0337fb5924ceb23de4648b8a81ae9d33b2b38/triton-3.6.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d002e07d7180fd65e622134fbd980c9a3d4211fb85224b56a0a0efbd422ab72f", size = 188400296, upload-time = "2026-01-20T16:00:56.042Z" }, +] + +[[package]] +name = "typeguard" +version = "4.5.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2b/e8/66e25efcc18542d58706ce4e50415710593721aae26e794ab1dec34fb66f/typeguard-4.5.1.tar.gz", hash = "sha256:f6f8ecbbc819c9bc749983cc67c02391e16a9b43b8b27f15dc70ed7c4a007274", size = 80121, upload-time = "2026-02-19T16:09:03.392Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/91/88/b55b3117287a8540b76dbdd87733808d4d01c8067a3b339408c250bb3600/typeguard-4.5.1-py3-none-any.whl", hash = "sha256:44d2bf329d49a244110a090b55f5f91aa82d9a9834ebfd30bcc73651e4a8cc40", size = 36745, upload-time = "2026-02-19T16:09:01.6Z" }, +] + +[[package]] +name = "typing-extensions" +version = "4.15.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, +] + +[[package]] +name = "typing-inspection" +version = "0.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/55/e3/70399cb7dd41c10ac53367ae42139cf4b1ca5f36bb3dc6c9d33acdb43655/typing_inspection-0.4.2.tar.gz", hash = "sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464", size = 75949, upload-time = "2025-10-01T02:14:41.687Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, +] + +[[package]] +name = "tyro" +version = "1.0.13" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "docstring-parser" }, + { name = "typeguard" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/24/d6/7126f9e7de139632134d59b5d1972e93c610ee2cb13829e8f4f48f6613cb/tyro-1.0.13.tar.gz", hash = "sha256:731a90c9836b77fffe7c3fa0477ef2d3b6fa91252ddc0bb4d32dadd4fcc143d4", size = 489479, upload-time = "2026-04-14T18:21:52.888Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/93/4f/c43a0a8f0c66fd40a1d6cc47332a5a1d1043e9b331f7070ea701b91a7598/tyro-1.0.13-py3-none-any.whl", hash = "sha256:a0bdb8462c551dd84fc00a76916ce4d37e879c84eefaf34e2165312407cc6c09", size = 185221, upload-time = "2026-04-14T18:21:54.328Z" }, +] + +[[package]] +name = "urllib3" +version = "2.6.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c7/24/5f1b3bdffd70275f6661c76461e25f024d5a38a46f04aaca912426a2b1d3/urllib3-2.6.3.tar.gz", hash = "sha256:1b62b6884944a57dbe321509ab94fd4d3b307075e0c2eae991ac71ee15ad38ed", size = 435556, upload-time = "2026-01-07T16:24:43.925Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl", hash = "sha256:bf272323e553dfb2e87d9bfd225ca7b0f467b919d7bbd355436d3fd37cb0acd4", size = 131584, upload-time = "2026-01-07T16:24:42.685Z" }, +] + +[[package]] +name = "vhacdx" +version = "0.0.10" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/15/e3/d2abc3dc4c1cb216c2efdc70b36f80efeb1bdbd7d420a676ddc9d9d980e1/vhacdx-0.0.10.tar.gz", hash = "sha256:fcc23201e319d79fe25e064847efc254bd39ac30af28cc761409e1f9142dd033", size = 58125, upload-time = "2025-12-02T20:58:45.358Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cf/9c/66375e65634c80f6efb46e81915126bf3e55dc9d6615217590cbc8316d2e/vhacdx-0.0.10-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7dd17d697d6d4d7cf66f1e947e0530041913981e05f7025236bec28a350b1a33", size = 224998, upload-time = "2025-12-02T20:57:51.639Z" }, + { url = "https://files.pythonhosted.org/packages/4e/e3/fc2644d3e7d0b2b52e2f681eb2878c0e1b9cafc53946f66736d0f01e237c/vhacdx-0.0.10-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:189ded39b709436cb732cdf694d4cf22e877aefb97e2ab2b55bf7ada9c030f93", size = 211130, upload-time = "2025-12-02T20:57:53.018Z" }, + { url = "https://files.pythonhosted.org/packages/e3/93/0b0f1977f5b3c2e1bbea5ef85e37a808ff73f1b7daf42950c57090e90dc7/vhacdx-0.0.10-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f3b03d35ab56a93beee338175dbe0b87552353e5dfb3ff37467e88f56cedf7cc", size = 239661, upload-time = "2025-12-02T20:57:54.144Z" }, + { url = "https://files.pythonhosted.org/packages/94/98/d2a6aeb1c6570a1fc1be29ee03db795f643ab03c6df7635522f23796b39d/vhacdx-0.0.10-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ea8c54ed193fa0db0248928fbf5d438b3872d615a506889d5b89fc6467d6411a", size = 252938, upload-time = "2025-12-02T20:57:55.275Z" }, + { url = "https://files.pythonhosted.org/packages/94/2e/1e678efc161a0d7fe1806f5e037ce11cc5964db7e08ccfc220ef63951863/vhacdx-0.0.10-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a5c898104140c72e4dc789e6125812671eee5e412916e83eff24a6148248ff5e", size = 1226696, upload-time = "2025-12-02T20:57:56.438Z" }, + { url = "https://files.pythonhosted.org/packages/90/5b/b302a0420a241c4910f4870eb9f39e6ada59858db441cc35bda511c17982/vhacdx-0.0.10-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:abdd0ba17786e578206594731df15c90e2751b6884220c8673124f47fd7ac620", size = 1287794, upload-time = "2025-12-02T20:57:57.694Z" }, + { url = "https://files.pythonhosted.org/packages/73/e9/f9729603ac75047a257f1b4ddac60cbde72b0abfd49ffed305751ba630a2/vhacdx-0.0.10-cp312-cp312-win_amd64.whl", hash = "sha256:79e7db59b4042295b21b79d55ba486a9a480550f696d466f158a30ed920dd0ec", size = 195033, upload-time = "2025-12-02T20:57:58.95Z" }, + { url = "https://files.pythonhosted.org/packages/0e/54/c2fc08d9324bbd92735caf9207cbabada3a8dd9d270d6e46ca69eb7f883d/vhacdx-0.0.10-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:0599bc2a96de8fc9aeff460b3e88b8572e84ae95b8fc6c9888ef4b92023c22d5", size = 225014, upload-time = "2025-12-02T20:58:00.938Z" }, + { url = "https://files.pythonhosted.org/packages/3b/9e/42adb642a12915acc9cb2bfab21710a6aabf045c26967ba0ff0e08a872d0/vhacdx-0.0.10-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:dc648829a1e973f34ee11393a4f334ab55e3e0e9c4b9f6d6349af966fdf1895a", size = 211127, upload-time = "2025-12-02T20:58:02.107Z" }, + { url = "https://files.pythonhosted.org/packages/51/3d/63e090cd966817b89643d7e52e13df45043b22a42c7fbf702866bdd75bc0/vhacdx-0.0.10-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:74c03f7315a434ec83cd0bff1e6bce6af4c01df61d677f48f3ffb36800606ee7", size = 239471, upload-time = "2025-12-02T20:58:03.173Z" }, + { url = "https://files.pythonhosted.org/packages/b8/b4/07ab1c828bae0eb5c72cd9a4cbe8b0376d374509be3c7055e1a399bf85c3/vhacdx-0.0.10-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1fcd02acc3733ec3a0a0d28ca7f526d4c56f14a3ceb4b12fce45acf72c09054a", size = 253019, upload-time = "2025-12-02T20:58:04.318Z" }, + { url = "https://files.pythonhosted.org/packages/05/cb/bc8a8858b300d2c092da11096ae0586ece446b4c41cb26620bf00d1d8232/vhacdx-0.0.10-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:4b9f8a80ca4c54d7fa76419a2ebd9e9386cd177dc4d2b97f2208ac57c9a7e8aa", size = 1226933, upload-time = "2025-12-02T20:58:05.907Z" }, + { url = "https://files.pythonhosted.org/packages/15/52/213230883874615f1661903bce1ace5013d03b34696efce8d53c662a3358/vhacdx-0.0.10-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:847bd43e82afb439dd3fa972618d786d0b98d8ef04a8e8a6381f6945204d2505", size = 1288871, upload-time = "2025-12-02T20:58:07.432Z" }, + { url = "https://files.pythonhosted.org/packages/32/25/f0e6806824f88d47ab8bc1c9bf6f11634fd7b382d635d0696825f3b5672f/vhacdx-0.0.10-cp313-cp313-win_amd64.whl", hash = "sha256:ab300c5f3fe4e54f99af92f9ea27c977b09df5f59190b0a3e025161110f71ce7", size = 195091, upload-time = "2025-12-02T20:58:08.783Z" }, +] + +[[package]] +name = "viser" +version = "1.0.26" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "imageio" }, + { name = "msgspec" }, + { name = "numpy" }, + { name = "requests" }, + { name = "rich" }, + { name = "tqdm" }, + { name = "trimesh" }, + { name = "typing-extensions" }, + { name = "websockets" }, + { name = "yourdfpy" }, + { name = "zstandard" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/12/ce/82a0e50fae21f5e02fcc5d9aff2ab59dccb9c319b6c4cf528f2228049b05/viser-1.0.26.tar.gz", hash = "sha256:dc08c6f505e70324b0603bdddf9714c00ac828c259ee49abd8ad094bfc90c91c", size = 4828261, upload-time = "2026-03-30T11:43:19.513Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/91/f7/762a2d5fab509d0c632b271e21e634462397cc02cca649771c3e9d2e0bcc/viser-1.0.26-py3-none-any.whl", hash = "sha256:03b177b4ef584f58f7b74fdf44cccb165b8a220ffd90728ef5c1e3d1b1fcf258", size = 4922888, upload-time = "2026-03-30T11:43:21.355Z" }, +] + +[[package]] +name = "wandb" +version = "0.26.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "gitpython" }, + { name = "packaging" }, + { name = "platformdirs" }, + { name = "protobuf" }, + { name = "pydantic" }, + { name = "pyyaml" }, + { name = "requests" }, + { name = "sentry-sdk" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/93/82/911948663ddf9e5ec6bc5cde19b0fffcb23c4b64a546bf5c084fde76c4cb/wandb-0.26.0.tar.gz", hash = "sha256:0356853895b53fe110e2ed17a1d49c15405498f08e5fbc339deab384f2df45f1", size = 42120837, upload-time = "2026-04-13T19:42:47.282Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/69/16/c0ea55323be74da9b4297d934b8e787251ae944d5776340c5498871927f1/wandb-0.26.0-py3-none-macosx_12_0_arm64.whl", hash = "sha256:1ece94a2a5eda1d0e3a2d8a2fd28aa0187705d6efa5ac4c0b8680083583b7ec1", size = 24800103, upload-time = "2026-04-13T19:42:23.946Z" }, + { url = "https://files.pythonhosted.org/packages/29/b8/4d38b43747616c4a9304be38b6e78526814deb5c1e01b3b6ebac82ce1cb5/wandb-0.26.0-py3-none-macosx_12_0_x86_64.whl", hash = "sha256:92f6f303fe2af50e3f711833a835150f9b4d8082874bfd9868cf15491ea2947e", size = 25956473, upload-time = "2026-04-13T19:42:26.772Z" }, + { url = "https://files.pythonhosted.org/packages/45/a6/940ca459d70c7cce7a6f7b395809f8ed051a25b2ce696fb93694c77f065f/wandb-0.26.0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:a17aae051a31831388cff880251c1b5bc38fbf6a283a0ee7c543709e8e9633d1", size = 25352442, upload-time = "2026-04-13T19:42:29.438Z" }, + { url = "https://files.pythonhosted.org/packages/4f/a8/55325da4b240d07ba2a8e1949a05b5942dd3346e14f7fd5e3cc72e46a648/wandb-0.26.0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:28a14ffc014e523498d077dfde12839b7be586ca8c3190e72e7167c1aea6ee4c", size = 27177821, upload-time = "2026-04-13T19:42:32.055Z" }, + { url = "https://files.pythonhosted.org/packages/c4/84/e4b0636a3e921e2cffb159b57b5a83787475993e2b5adb6181fbf7712a59/wandb-0.26.0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:fb9a63babeee044fecf65a4675f7dfb0efaea4986e498a3bc8f948558af877e7", size = 25522688, upload-time = "2026-04-13T19:42:34.625Z" }, + { url = "https://files.pythonhosted.org/packages/dd/f0/821a451110dd5f5c39358752abbdcb56c4fcebcc41039c7dcd4b024d2e27/wandb-0.26.0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:3eb88fb556a64bf4492cf571bb851d47871901c096f0540d841ccb50f5cbeb66", size = 27486467, upload-time = "2026-04-13T19:42:37.119Z" }, + { url = "https://files.pythonhosted.org/packages/24/c5/a4eb8fb6e7527584c6ccdf5c9b265283ed0c9d94d26d5eb28f9b48cd5779/wandb-0.26.0-py3-none-win32.whl", hash = "sha256:362828d48d21dd4877e28fdce40421ebdfc16d1fe0b59e8371b12d75bbc3f1e7", size = 24908555, upload-time = "2026-04-13T19:42:39.416Z" }, + { url = "https://files.pythonhosted.org/packages/35/3d/bf182f3af977e6297fc05bc3fd9bd51feacfe4d2c4ce83c90eb2ad7ce59b/wandb-0.26.0-py3-none-win_amd64.whl", hash = "sha256:21a8346434fd30e1bc13a26b226fc29b6f33a1cb346d610cbcb4040c3b0e1f63", size = 24908559, upload-time = "2026-04-13T19:42:42.019Z" }, + { url = "https://files.pythonhosted.org/packages/1a/3e/344cb29b593f8e7abc14cc268dafde1974bae3f073b4885476f4fbba3cb8/wandb-0.26.0-py3-none-win_arm64.whl", hash = "sha256:99bd11974e9005d3a3f82e1fabfc4909ffa1fdede23a8839f5fbaea2f5be9033", size = 22936140, upload-time = "2026-04-13T19:42:44.715Z" }, +] + +[[package]] +name = "warp-lang" +version = "1.13.0.dev20260415" +source = { registry = "https://pypi.nvidia.com/" } +dependencies = [ + { name = "numpy" }, +] +wheels = [ + { url = "https://pypi.nvidia.cn/warp-lang/warp_lang-1.13.0.dev20260415-py3-none-macosx_11_0_arm64.whl", hash = "sha256:a9bde7f41bfd2e89094672b5d3cf1adefc1fe420115348016fd9299cd6778866" }, + { url = "https://pypi.nvidia.cn/warp-lang/warp_lang-1.13.0.dev20260415-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:dd409f3cb1c19912eaa8a1d9756a635ca83ee1c7cd3679d550d3b501a7460784" }, + { url = "https://pypi.nvidia.cn/warp-lang/warp_lang-1.13.0.dev20260415-py3-none-manylinux_2_34_aarch64.whl", hash = "sha256:7bb1d4c39beb94af644097e122c3c45b114d2a2d8119fbc19de0521ad3b748d4" }, + { url = "https://pypi.nvidia.cn/warp-lang/warp_lang-1.13.0.dev20260415-py3-none-win_amd64.whl", hash = "sha256:ec3fa2e664b908367069e7e36bee23501d48dcb8900ed90422487ba1fc52178e" }, +] + +[[package]] +name = "wcwidth" +version = "0.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/35/a2/8e3becb46433538a38726c948d3399905a4c7cabd0df578ede5dc51f0ec2/wcwidth-0.6.0.tar.gz", hash = "sha256:cdc4e4262d6ef9a1a57e018384cbeb1208d8abbc64176027e2c2455c81313159", size = 159684, upload-time = "2026-02-06T19:19:40.919Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/68/5a/199c59e0a824a3db2b89c5d2dade7ab5f9624dbf6448dc291b46d5ec94d3/wcwidth-0.6.0-py3-none-any.whl", hash = "sha256:1a3a1e510b553315f8e146c54764f4fb6264ffad731b3d78088cdb1478ffbdad", size = 94189, upload-time = "2026-02-06T19:19:39.646Z" }, +] + +[[package]] +name = "websockets" +version = "16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/04/24/4b2031d72e840ce4c1ccb255f693b15c334757fc50023e4db9537080b8c4/websockets-16.0.tar.gz", hash = "sha256:5f6261a5e56e8d5c42a4497b364ea24d94d9563e8fbd44e78ac40879c60179b5", size = 179346, upload-time = "2026-01-10T09:23:47.181Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/84/7b/bac442e6b96c9d25092695578dda82403c77936104b5682307bd4deb1ad4/websockets-16.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:71c989cbf3254fbd5e84d3bff31e4da39c43f884e64f2551d14bb3c186230f00", size = 177365, upload-time = "2026-01-10T09:22:46.787Z" }, + { url = "https://files.pythonhosted.org/packages/b0/fe/136ccece61bd690d9c1f715baaeefd953bb2360134de73519d5df19d29ca/websockets-16.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:8b6e209ffee39ff1b6d0fa7bfef6de950c60dfb91b8fcead17da4ee539121a79", size = 175038, upload-time = "2026-01-10T09:22:47.999Z" }, + { url = "https://files.pythonhosted.org/packages/40/1e/9771421ac2286eaab95b8575b0cb701ae3663abf8b5e1f64f1fd90d0a673/websockets-16.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:86890e837d61574c92a97496d590968b23c2ef0aeb8a9bc9421d174cd378ae39", size = 175328, upload-time = "2026-01-10T09:22:49.809Z" }, + { url = "https://files.pythonhosted.org/packages/18/29/71729b4671f21e1eaa5d6573031ab810ad2936c8175f03f97f3ff164c802/websockets-16.0-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:9b5aca38b67492ef518a8ab76851862488a478602229112c4b0d58d63a7a4d5c", size = 184915, upload-time = "2026-01-10T09:22:51.071Z" }, + { url = "https://files.pythonhosted.org/packages/97/bb/21c36b7dbbafc85d2d480cd65df02a1dc93bf76d97147605a8e27ff9409d/websockets-16.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e0334872c0a37b606418ac52f6ab9cfd17317ac26365f7f65e203e2d0d0d359f", size = 186152, upload-time = "2026-01-10T09:22:52.224Z" }, + { url = "https://files.pythonhosted.org/packages/4a/34/9bf8df0c0cf88fa7bfe36678dc7b02970c9a7d5e065a3099292db87b1be2/websockets-16.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a0b31e0b424cc6b5a04b8838bbaec1688834b2383256688cf47eb97412531da1", size = 185583, upload-time = "2026-01-10T09:22:53.443Z" }, + { url = "https://files.pythonhosted.org/packages/47/88/4dd516068e1a3d6ab3c7c183288404cd424a9a02d585efbac226cb61ff2d/websockets-16.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:485c49116d0af10ac698623c513c1cc01c9446c058a4e61e3bf6c19dff7335a2", size = 184880, upload-time = "2026-01-10T09:22:55.033Z" }, + { url = "https://files.pythonhosted.org/packages/91/d6/7d4553ad4bf1c0421e1ebd4b18de5d9098383b5caa1d937b63df8d04b565/websockets-16.0-cp312-cp312-win32.whl", hash = "sha256:eaded469f5e5b7294e2bdca0ab06becb6756ea86894a47806456089298813c89", size = 178261, upload-time = "2026-01-10T09:22:56.251Z" }, + { url = "https://files.pythonhosted.org/packages/c3/f0/f3a17365441ed1c27f850a80b2bc680a0fa9505d733fe152fdf5e98c1c0b/websockets-16.0-cp312-cp312-win_amd64.whl", hash = "sha256:5569417dc80977fc8c2d43a86f78e0a5a22fee17565d78621b6bb264a115d4ea", size = 178693, upload-time = "2026-01-10T09:22:57.478Z" }, + { url = "https://files.pythonhosted.org/packages/cc/9c/baa8456050d1c1b08dd0ec7346026668cbc6f145ab4e314d707bb845bf0d/websockets-16.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:878b336ac47938b474c8f982ac2f7266a540adc3fa4ad74ae96fea9823a02cc9", size = 177364, upload-time = "2026-01-10T09:22:59.333Z" }, + { url = "https://files.pythonhosted.org/packages/7e/0c/8811fc53e9bcff68fe7de2bcbe75116a8d959ac699a3200f4847a8925210/websockets-16.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:52a0fec0e6c8d9a784c2c78276a48a2bdf099e4ccc2a4cad53b27718dbfd0230", size = 175039, upload-time = "2026-01-10T09:23:01.171Z" }, + { url = "https://files.pythonhosted.org/packages/aa/82/39a5f910cb99ec0b59e482971238c845af9220d3ab9fa76dd9162cda9d62/websockets-16.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:e6578ed5b6981005df1860a56e3617f14a6c307e6a71b4fff8c48fdc50f3ed2c", size = 175323, upload-time = "2026-01-10T09:23:02.341Z" }, + { url = "https://files.pythonhosted.org/packages/bd/28/0a25ee5342eb5d5f297d992a77e56892ecb65e7854c7898fb7d35e9b33bd/websockets-16.0-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:95724e638f0f9c350bb1c2b0a7ad0e83d9cc0c9259f3ea94e40d7b02a2179ae5", size = 184975, upload-time = "2026-01-10T09:23:03.756Z" }, + { url = "https://files.pythonhosted.org/packages/f9/66/27ea52741752f5107c2e41fda05e8395a682a1e11c4e592a809a90c6a506/websockets-16.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c0204dc62a89dc9d50d682412c10b3542d748260d743500a85c13cd1ee4bde82", size = 186203, upload-time = "2026-01-10T09:23:05.01Z" }, + { url = "https://files.pythonhosted.org/packages/37/e5/8e32857371406a757816a2b471939d51c463509be73fa538216ea52b792a/websockets-16.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:52ac480f44d32970d66763115edea932f1c5b1312de36df06d6b219f6741eed8", size = 185653, upload-time = "2026-01-10T09:23:06.301Z" }, + { url = "https://files.pythonhosted.org/packages/9b/67/f926bac29882894669368dc73f4da900fcdf47955d0a0185d60103df5737/websockets-16.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6e5a82b677f8f6f59e8dfc34ec06ca6b5b48bc4fcda346acd093694cc2c24d8f", size = 184920, upload-time = "2026-01-10T09:23:07.492Z" }, + { url = "https://files.pythonhosted.org/packages/3c/a1/3d6ccdcd125b0a42a311bcd15a7f705d688f73b2a22d8cf1c0875d35d34a/websockets-16.0-cp313-cp313-win32.whl", hash = "sha256:abf050a199613f64c886ea10f38b47770a65154dc37181bfaff70c160f45315a", size = 178255, upload-time = "2026-01-10T09:23:09.245Z" }, + { url = "https://files.pythonhosted.org/packages/6b/ae/90366304d7c2ce80f9b826096a9e9048b4bb760e44d3b873bb272cba696b/websockets-16.0-cp313-cp313-win_amd64.whl", hash = "sha256:3425ac5cf448801335d6fdc7ae1eb22072055417a96cc6b31b3861f455fbc156", size = 178689, upload-time = "2026-01-10T09:23:10.483Z" }, + { url = "https://files.pythonhosted.org/packages/6f/28/258ebab549c2bf3e64d2b0217b973467394a9cea8c42f70418ca2c5d0d2e/websockets-16.0-py3-none-any.whl", hash = "sha256:1637db62fad1dc833276dded54215f2c7fa46912301a24bd94d45d46a011ceec", size = 171598, upload-time = "2026-01-10T09:23:45.395Z" }, +] + +[[package]] +name = "werkzeug" +version = "3.1.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/dd/b2/381be8cfdee792dd117872481b6e378f85c957dd7c5bca38897b08f765fd/werkzeug-3.1.8.tar.gz", hash = "sha256:9bad61a4268dac112f1c5cd4630a56ede601b6ed420300677a869083d70a4c44", size = 875852, upload-time = "2026-04-02T18:49:14.268Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/93/8c/2e650f2afeb7ee576912636c23ddb621c91ac6a98e66dc8d29c3c69446e1/werkzeug-3.1.8-py3-none-any.whl", hash = "sha256:63a77fb8892bf28ebc3178683445222aa500e48ebad5ec77b0ad80f8726b1f50", size = 226459, upload-time = "2026-04-02T18:49:12.72Z" }, +] + +[[package]] +name = "wrapt" +version = "2.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2e/64/925f213fdcbb9baeb1530449ac71a4d57fc361c053d06bf78d0c5c7cd80c/wrapt-2.1.2.tar.gz", hash = "sha256:3996a67eecc2c68fd47b4e3c564405a5777367adfd9b8abb58387b63ee83b21e", size = 81678, upload-time = "2026-03-06T02:53:25.134Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4c/b6/1db817582c49c7fcbb7df6809d0f515af29d7c2fbf57eb44c36e98fb1492/wrapt-2.1.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:ff2aad9c4cda28a8f0653fc2d487596458c2a3f475e56ba02909e950a9efa6a9", size = 61255, upload-time = "2026-03-06T02:52:45.663Z" }, + { url = "https://files.pythonhosted.org/packages/a2/16/9b02a6b99c09227c93cd4b73acc3678114154ec38da53043c0ddc1fba0dc/wrapt-2.1.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6433ea84e1cfacf32021d2a4ee909554ade7fd392caa6f7c13f1f4bf7b8e8748", size = 61848, upload-time = "2026-03-06T02:53:48.728Z" }, + { url = "https://files.pythonhosted.org/packages/af/aa/ead46a88f9ec3a432a4832dfedb84092fc35af2d0ba40cd04aea3889f247/wrapt-2.1.2-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:c20b757c268d30d6215916a5fa8461048d023865d888e437fab451139cad6c8e", size = 121433, upload-time = "2026-03-06T02:54:40.328Z" }, + { url = "https://files.pythonhosted.org/packages/3a/9f/742c7c7cdf58b59085a1ee4b6c37b013f66ac33673a7ef4aaed5e992bc33/wrapt-2.1.2-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:79847b83eb38e70d93dc392c7c5b587efe65b3e7afcc167aa8abd5d60e8761c8", size = 123013, upload-time = "2026-03-06T02:53:26.58Z" }, + { url = "https://files.pythonhosted.org/packages/e8/44/2c3dd45d53236b7ed7c646fcf212251dc19e48e599debd3926b52310fafb/wrapt-2.1.2-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f8fba1bae256186a83d1875b2b1f4e2d1242e8fac0f58ec0d7e41b26967b965c", size = 117326, upload-time = "2026-03-06T02:53:11.547Z" }, + { url = "https://files.pythonhosted.org/packages/74/e2/b17d66abc26bd96f89dec0ecd0ef03da4a1286e6ff793839ec431b9fae57/wrapt-2.1.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e3d3b35eedcf5f7d022291ecd7533321c4775f7b9cd0050a31a68499ba45757c", size = 121444, upload-time = "2026-03-06T02:54:09.5Z" }, + { url = "https://files.pythonhosted.org/packages/3c/62/e2977843fdf9f03daf1586a0ff49060b1b2fc7ff85a7ea82b6217c1ae36e/wrapt-2.1.2-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:6f2c5390460de57fa9582bc8a1b7a6c86e1a41dfad74c5225fc07044c15cc8d1", size = 116237, upload-time = "2026-03-06T02:54:03.884Z" }, + { url = "https://files.pythonhosted.org/packages/88/dd/27fc67914e68d740bce512f11734aec08696e6b17641fef8867c00c949fc/wrapt-2.1.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:7dfa9f2cf65d027b951d05c662cc99ee3bd01f6e4691ed39848a7a5fffc902b2", size = 120563, upload-time = "2026-03-06T02:53:20.412Z" }, + { url = "https://files.pythonhosted.org/packages/ec/9f/b750b3692ed2ef4705cb305bd68858e73010492b80e43d2a4faa5573cbe7/wrapt-2.1.2-cp312-cp312-win32.whl", hash = "sha256:eba8155747eb2cae4a0b913d9ebd12a1db4d860fc4c829d7578c7b989bd3f2f0", size = 58198, upload-time = "2026-03-06T02:53:37.732Z" }, + { url = "https://files.pythonhosted.org/packages/8e/b2/feecfe29f28483d888d76a48f03c4c4d8afea944dbee2b0cd3380f9df032/wrapt-2.1.2-cp312-cp312-win_amd64.whl", hash = "sha256:1c51c738d7d9faa0b3601708e7e2eda9bf779e1b601dce6c77411f2a1b324a63", size = 60441, upload-time = "2026-03-06T02:52:47.138Z" }, + { url = "https://files.pythonhosted.org/packages/44/e1/e328f605d6e208547ea9fd120804fcdec68536ac748987a68c47c606eea8/wrapt-2.1.2-cp312-cp312-win_arm64.whl", hash = "sha256:c8e46ae8e4032792eb2f677dbd0d557170a8e5524d22acc55199f43efedd39bf", size = 58836, upload-time = "2026-03-06T02:53:22.053Z" }, + { url = "https://files.pythonhosted.org/packages/4c/7a/d936840735c828b38d26a854e85d5338894cda544cb7a85a9d5b8b9c4df7/wrapt-2.1.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:787fd6f4d67befa6fe2abdffcbd3de2d82dfc6fb8a6d850407c53332709d030b", size = 61259, upload-time = "2026-03-06T02:53:41.922Z" }, + { url = "https://files.pythonhosted.org/packages/5e/88/9a9b9a90ac8ca11c2fdb6a286cb3a1fc7dd774c00ed70929a6434f6bc634/wrapt-2.1.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:4bdf26e03e6d0da3f0e9422fd36bcebf7bc0eeb55fdf9c727a09abc6b9fe472e", size = 61851, upload-time = "2026-03-06T02:52:48.672Z" }, + { url = "https://files.pythonhosted.org/packages/03/a9/5b7d6a16fd6533fed2756900fc8fc923f678179aea62ada6d65c92718c00/wrapt-2.1.2-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:bbac24d879aa22998e87f6b3f481a5216311e7d53c7db87f189a7a0266dafffb", size = 121446, upload-time = "2026-03-06T02:54:14.013Z" }, + { url = "https://files.pythonhosted.org/packages/45/bb/34c443690c847835cfe9f892be78c533d4f32366ad2888972c094a897e39/wrapt-2.1.2-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:16997dfb9d67addc2e3f41b62a104341e80cac52f91110dece393923c0ebd5ca", size = 123056, upload-time = "2026-03-06T02:54:10.829Z" }, + { url = "https://files.pythonhosted.org/packages/93/b9/ff205f391cb708f67f41ea148545f2b53ff543a7ac293b30d178af4d2271/wrapt-2.1.2-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:162e4e2ba7542da9027821cb6e7c5e068d64f9a10b5f15512ea28e954893a267", size = 117359, upload-time = "2026-03-06T02:53:03.623Z" }, + { url = "https://files.pythonhosted.org/packages/1f/3d/1ea04d7747825119c3c9a5e0874a40b33594ada92e5649347c457d982805/wrapt-2.1.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f29c827a8d9936ac320746747a016c4bc66ef639f5cd0d32df24f5eacbf9c69f", size = 121479, upload-time = "2026-03-06T02:53:45.844Z" }, + { url = "https://files.pythonhosted.org/packages/78/cc/ee3a011920c7a023b25e8df26f306b2484a531ab84ca5c96260a73de76c0/wrapt-2.1.2-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:a9dd9813825f7ecb018c17fd147a01845eb330254dff86d3b5816f20f4d6aaf8", size = 116271, upload-time = "2026-03-06T02:54:46.356Z" }, + { url = "https://files.pythonhosted.org/packages/98/fd/e5ff7ded41b76d802cf1191288473e850d24ba2e39a6ec540f21ae3b57cb/wrapt-2.1.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6f8dbdd3719e534860d6a78526aafc220e0241f981367018c2875178cf83a413", size = 120573, upload-time = "2026-03-06T02:52:50.163Z" }, + { url = "https://files.pythonhosted.org/packages/47/c5/242cae3b5b080cd09bacef0591691ba1879739050cc7c801ff35c8886b66/wrapt-2.1.2-cp313-cp313-win32.whl", hash = "sha256:5c35b5d82b16a3bc6e0a04349b606a0582bc29f573786aebe98e0c159bc48db6", size = 58205, upload-time = "2026-03-06T02:53:47.494Z" }, + { url = "https://files.pythonhosted.org/packages/12/69/c358c61e7a50f290958809b3c61ebe8b3838ea3e070d7aac9814f95a0528/wrapt-2.1.2-cp313-cp313-win_amd64.whl", hash = "sha256:f8bc1c264d8d1cf5b3560a87bbdd31131573eb25f9f9447bb6252b8d4c44a3a1", size = 60452, upload-time = "2026-03-06T02:53:30.038Z" }, + { url = "https://files.pythonhosted.org/packages/8e/66/c8a6fcfe321295fd8c0ab1bd685b5a01462a9b3aa2f597254462fc2bc975/wrapt-2.1.2-cp313-cp313-win_arm64.whl", hash = "sha256:3beb22f674550d5634642c645aba4c72a2c66fb185ae1aebe1e955fae5a13baf", size = 58842, upload-time = "2026-03-06T02:52:52.114Z" }, + { url = "https://files.pythonhosted.org/packages/da/55/9c7052c349106e0b3f17ae8db4b23a691a963c334de7f9dbd60f8f74a831/wrapt-2.1.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0fc04bc8664a8bc4c8e00b37b5355cffca2535209fba1abb09ae2b7c76ddf82b", size = 63075, upload-time = "2026-03-06T02:53:19.108Z" }, + { url = "https://files.pythonhosted.org/packages/09/a8/ce7b4006f7218248dd71b7b2b732d0710845a0e49213b18faef64811ffef/wrapt-2.1.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:a9b9d50c9af998875a1482a038eb05755dfd6fe303a313f6a940bb53a83c3f18", size = 63719, upload-time = "2026-03-06T02:54:33.452Z" }, + { url = "https://files.pythonhosted.org/packages/e4/e5/2ca472e80b9e2b7a17f106bb8f9df1db11e62101652ce210f66935c6af67/wrapt-2.1.2-cp313-cp313t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:2d3ff4f0024dd224290c0eabf0240f1bfc1f26363431505fb1b0283d3b08f11d", size = 152643, upload-time = "2026-03-06T02:52:42.721Z" }, + { url = "https://files.pythonhosted.org/packages/36/42/30f0f2cefca9d9cbf6835f544d825064570203c3e70aa873d8ae12e23791/wrapt-2.1.2-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3278c471f4468ad544a691b31bb856374fbdefb7fee1a152153e64019379f015", size = 158805, upload-time = "2026-03-06T02:54:25.441Z" }, + { url = "https://files.pythonhosted.org/packages/bb/67/d08672f801f604889dcf58f1a0b424fe3808860ede9e03affc1876b295af/wrapt-2.1.2-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a8914c754d3134a3032601c6984db1c576e6abaf3fc68094bb8ab1379d75ff92", size = 145990, upload-time = "2026-03-06T02:53:57.456Z" }, + { url = "https://files.pythonhosted.org/packages/68/a7/fd371b02e73babec1de6ade596e8cd9691051058cfdadbfd62a5898f3295/wrapt-2.1.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:ff95d4264e55839be37bafe1536db2ab2de19da6b65f9244f01f332b5286cfbf", size = 155670, upload-time = "2026-03-06T02:54:55.309Z" }, + { url = "https://files.pythonhosted.org/packages/86/2d/9fe0095dfdb621009f40117dcebf41d7396c2c22dca6eac779f4c007b86c/wrapt-2.1.2-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:76405518ca4e1b76fbb1b9f686cff93aebae03920cc55ceeec48ff9f719c5f67", size = 144357, upload-time = "2026-03-06T02:54:24.092Z" }, + { url = "https://files.pythonhosted.org/packages/0e/b6/ec7b4a254abbe4cde9fa15c5d2cca4518f6b07d0f1b77d4ee9655e30280e/wrapt-2.1.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:c0be8b5a74c5824e9359b53e7e58bef71a729bacc82e16587db1c4ebc91f7c5a", size = 150269, upload-time = "2026-03-06T02:53:31.268Z" }, + { url = "https://files.pythonhosted.org/packages/6e/6b/2fabe8ebf148f4ee3c782aae86a795cc68ffe7d432ef550f234025ce0cfa/wrapt-2.1.2-cp313-cp313t-win32.whl", hash = "sha256:f01277d9a5fc1862f26f7626da9cf443bebc0abd2f303f41c5e995b15887dabd", size = 59894, upload-time = "2026-03-06T02:54:15.391Z" }, + { url = "https://files.pythonhosted.org/packages/ca/fb/9ba66fc2dedc936de5f8073c0217b5d4484e966d87723415cc8262c5d9c2/wrapt-2.1.2-cp313-cp313t-win_amd64.whl", hash = "sha256:84ce8f1c2104d2f6daa912b1b5b039f331febfeee74f8042ad4e04992bd95c8f", size = 63197, upload-time = "2026-03-06T02:54:41.943Z" }, + { url = "https://files.pythonhosted.org/packages/c0/1c/012d7423c95d0e337117723eb8ecf73c622ce15a97847e84cf3f8f26cd7e/wrapt-2.1.2-cp313-cp313t-win_arm64.whl", hash = "sha256:a93cd767e37faeddbe07d8fc4212d5cba660af59bdb0f6372c93faaa13e6e679", size = 60363, upload-time = "2026-03-06T02:54:48.093Z" }, + { url = "https://files.pythonhosted.org/packages/1a/c7/8528ac2dfa2c1e6708f647df7ae144ead13f0a31146f43c7264b4942bf12/wrapt-2.1.2-py3-none-any.whl", hash = "sha256:b8fd6fa2b2c4e7621808f8c62e8317f4aae56e59721ad933bac5239d913cf0e8", size = 43993, upload-time = "2026-03-06T02:53:12.905Z" }, +] + +[[package]] +name = "xxhash" +version = "3.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/02/84/30869e01909fb37a6cc7e18688ee8bf1e42d57e7e0777636bd47524c43c7/xxhash-3.6.0.tar.gz", hash = "sha256:f0162a78b13a0d7617b2845b90c763339d1f1d82bb04a4b07f4ab535cc5e05d6", size = 85160, upload-time = "2025-10-02T14:37:08.097Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9a/07/d9412f3d7d462347e4511181dea65e47e0d0e16e26fbee2ea86a2aefb657/xxhash-3.6.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:01362c4331775398e7bb34e3ab403bc9ee9f7c497bc7dee6272114055277dd3c", size = 32744, upload-time = "2025-10-02T14:34:34.622Z" }, + { url = "https://files.pythonhosted.org/packages/79/35/0429ee11d035fc33abe32dca1b2b69e8c18d236547b9a9b72c1929189b9a/xxhash-3.6.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b7b2df81a23f8cb99656378e72501b2cb41b1827c0f5a86f87d6b06b69f9f204", size = 30816, upload-time = "2025-10-02T14:34:36.043Z" }, + { url = "https://files.pythonhosted.org/packages/b7/f2/57eb99aa0f7d98624c0932c5b9a170e1806406cdbcdb510546634a1359e0/xxhash-3.6.0-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:dc94790144e66b14f67b10ac8ed75b39ca47536bf8800eb7c24b50271ea0c490", size = 194035, upload-time = "2025-10-02T14:34:37.354Z" }, + { url = "https://files.pythonhosted.org/packages/4c/ed/6224ba353690d73af7a3f1c7cdb1fc1b002e38f783cb991ae338e1eb3d79/xxhash-3.6.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:93f107c673bccf0d592cdba077dedaf52fe7f42dcd7676eba1f6d6f0c3efffd2", size = 212914, upload-time = "2025-10-02T14:34:38.6Z" }, + { url = "https://files.pythonhosted.org/packages/38/86/fb6b6130d8dd6b8942cc17ab4d90e223653a89aa32ad2776f8af7064ed13/xxhash-3.6.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:2aa5ee3444c25b69813663c9f8067dcfaa2e126dc55e8dddf40f4d1c25d7effa", size = 212163, upload-time = "2025-10-02T14:34:39.872Z" }, + { url = "https://files.pythonhosted.org/packages/ee/dc/e84875682b0593e884ad73b2d40767b5790d417bde603cceb6878901d647/xxhash-3.6.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f7f99123f0e1194fa59cc69ad46dbae2e07becec5df50a0509a808f90a0f03f0", size = 445411, upload-time = "2025-10-02T14:34:41.569Z" }, + { url = "https://files.pythonhosted.org/packages/11/4f/426f91b96701ec2f37bb2b8cec664eff4f658a11f3fa9d94f0a887ea6d2b/xxhash-3.6.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:49e03e6fe2cac4a1bc64952dd250cf0dbc5ef4ebb7b8d96bce82e2de163c82a2", size = 193883, upload-time = "2025-10-02T14:34:43.249Z" }, + { url = "https://files.pythonhosted.org/packages/53/5a/ddbb83eee8e28b778eacfc5a85c969673e4023cdeedcfcef61f36731610b/xxhash-3.6.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:bd17fede52a17a4f9a7bc4472a5867cb0b160deeb431795c0e4abe158bc784e9", size = 210392, upload-time = "2025-10-02T14:34:45.042Z" }, + { url = "https://files.pythonhosted.org/packages/1e/c2/ff69efd07c8c074ccdf0a4f36fcdd3d27363665bcdf4ba399abebe643465/xxhash-3.6.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:6fb5f5476bef678f69db04f2bd1efbed3030d2aba305b0fc1773645f187d6a4e", size = 197898, upload-time = "2025-10-02T14:34:46.302Z" }, + { url = "https://files.pythonhosted.org/packages/58/ca/faa05ac19b3b622c7c9317ac3e23954187516298a091eb02c976d0d3dd45/xxhash-3.6.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:843b52f6d88071f87eba1631b684fcb4b2068cd2180a0224122fe4ef011a9374", size = 210655, upload-time = "2025-10-02T14:34:47.571Z" }, + { url = "https://files.pythonhosted.org/packages/d4/7a/06aa7482345480cc0cb597f5c875b11a82c3953f534394f620b0be2f700c/xxhash-3.6.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:7d14a6cfaf03b1b6f5f9790f76880601ccc7896aff7ab9cd8978a939c1eb7e0d", size = 414001, upload-time = "2025-10-02T14:34:49.273Z" }, + { url = "https://files.pythonhosted.org/packages/23/07/63ffb386cd47029aa2916b3d2f454e6cc5b9f5c5ada3790377d5430084e7/xxhash-3.6.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:418daf3db71e1413cfe211c2f9a528456936645c17f46b5204705581a45390ae", size = 191431, upload-time = "2025-10-02T14:34:50.798Z" }, + { url = "https://files.pythonhosted.org/packages/0f/93/14fde614cadb4ddf5e7cebf8918b7e8fac5ae7861c1875964f17e678205c/xxhash-3.6.0-cp312-cp312-win32.whl", hash = "sha256:50fc255f39428a27299c20e280d6193d8b63b8ef8028995323bf834a026b4fbb", size = 30617, upload-time = "2025-10-02T14:34:51.954Z" }, + { url = "https://files.pythonhosted.org/packages/13/5d/0d125536cbe7565a83d06e43783389ecae0c0f2ed037b48ede185de477c0/xxhash-3.6.0-cp312-cp312-win_amd64.whl", hash = "sha256:c0f2ab8c715630565ab8991b536ecded9416d615538be8ecddce43ccf26cbc7c", size = 31534, upload-time = "2025-10-02T14:34:53.276Z" }, + { url = "https://files.pythonhosted.org/packages/54/85/6ec269b0952ec7e36ba019125982cf11d91256a778c7c3f98a4c5043d283/xxhash-3.6.0-cp312-cp312-win_arm64.whl", hash = "sha256:eae5c13f3bc455a3bbb68bdc513912dc7356de7e2280363ea235f71f54064829", size = 27876, upload-time = "2025-10-02T14:34:54.371Z" }, + { url = "https://files.pythonhosted.org/packages/33/76/35d05267ac82f53ae9b0e554da7c5e281ee61f3cad44c743f0fcd354f211/xxhash-3.6.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:599e64ba7f67472481ceb6ee80fa3bd828fd61ba59fb11475572cc5ee52b89ec", size = 32738, upload-time = "2025-10-02T14:34:55.839Z" }, + { url = "https://files.pythonhosted.org/packages/31/a8/3fbce1cd96534a95e35d5120637bf29b0d7f5d8fa2f6374e31b4156dd419/xxhash-3.6.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:7d8b8aaa30fca4f16f0c84a5c8d7ddee0e25250ec2796c973775373257dde8f1", size = 30821, upload-time = "2025-10-02T14:34:57.219Z" }, + { url = "https://files.pythonhosted.org/packages/0c/ea/d387530ca7ecfa183cb358027f1833297c6ac6098223fd14f9782cd0015c/xxhash-3.6.0-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:d597acf8506d6e7101a4a44a5e428977a51c0fadbbfd3c39650cca9253f6e5a6", size = 194127, upload-time = "2025-10-02T14:34:59.21Z" }, + { url = "https://files.pythonhosted.org/packages/ba/0c/71435dcb99874b09a43b8d7c54071e600a7481e42b3e3ce1eb5226a5711a/xxhash-3.6.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:858dc935963a33bc33490128edc1c12b0c14d9c7ebaa4e387a7869ecc4f3e263", size = 212975, upload-time = "2025-10-02T14:35:00.816Z" }, + { url = "https://files.pythonhosted.org/packages/84/7a/c2b3d071e4bb4a90b7057228a99b10d51744878f4a8a6dd643c8bd897620/xxhash-3.6.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ba284920194615cb8edf73bf52236ce2e1664ccd4a38fdb543506413529cc546", size = 212241, upload-time = "2025-10-02T14:35:02.207Z" }, + { url = "https://files.pythonhosted.org/packages/81/5f/640b6eac0128e215f177df99eadcd0f1b7c42c274ab6a394a05059694c5a/xxhash-3.6.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:4b54219177f6c6674d5378bd862c6aedf64725f70dd29c472eaae154df1a2e89", size = 445471, upload-time = "2025-10-02T14:35:03.61Z" }, + { url = "https://files.pythonhosted.org/packages/5e/1e/3c3d3ef071b051cc3abbe3721ffb8365033a172613c04af2da89d5548a87/xxhash-3.6.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:42c36dd7dbad2f5238950c377fcbf6811b1cdb1c444fab447960030cea60504d", size = 193936, upload-time = "2025-10-02T14:35:05.013Z" }, + { url = "https://files.pythonhosted.org/packages/2c/bd/4a5f68381939219abfe1c22a9e3a5854a4f6f6f3c4983a87d255f21f2e5d/xxhash-3.6.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f22927652cba98c44639ffdc7aaf35828dccf679b10b31c4ad72a5b530a18eb7", size = 210440, upload-time = "2025-10-02T14:35:06.239Z" }, + { url = "https://files.pythonhosted.org/packages/eb/37/b80fe3d5cfb9faff01a02121a0f4d565eb7237e9e5fc66e73017e74dcd36/xxhash-3.6.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:b45fad44d9c5c119e9c6fbf2e1c656a46dc68e280275007bbfd3d572b21426db", size = 197990, upload-time = "2025-10-02T14:35:07.735Z" }, + { url = "https://files.pythonhosted.org/packages/d7/fd/2c0a00c97b9e18f72e1f240ad4e8f8a90fd9d408289ba9c7c495ed7dc05c/xxhash-3.6.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:6f2580ffab1a8b68ef2b901cde7e55fa8da5e4be0977c68f78fc80f3c143de42", size = 210689, upload-time = "2025-10-02T14:35:09.438Z" }, + { url = "https://files.pythonhosted.org/packages/93/86/5dd8076a926b9a95db3206aba20d89a7fc14dd5aac16e5c4de4b56033140/xxhash-3.6.0-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:40c391dd3cd041ebc3ffe6f2c862f402e306eb571422e0aa918d8070ba31da11", size = 414068, upload-time = "2025-10-02T14:35:11.162Z" }, + { url = "https://files.pythonhosted.org/packages/af/3c/0bb129170ee8f3650f08e993baee550a09593462a5cddd8e44d0011102b1/xxhash-3.6.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:f205badabde7aafd1a31e8ca2a3e5a763107a71c397c4481d6a804eb5063d8bd", size = 191495, upload-time = "2025-10-02T14:35:12.971Z" }, + { url = "https://files.pythonhosted.org/packages/e9/3a/6797e0114c21d1725e2577508e24006fd7ff1d8c0c502d3b52e45c1771d8/xxhash-3.6.0-cp313-cp313-win32.whl", hash = "sha256:2577b276e060b73b73a53042ea5bd5203d3e6347ce0d09f98500f418a9fcf799", size = 30620, upload-time = "2025-10-02T14:35:14.129Z" }, + { url = "https://files.pythonhosted.org/packages/86/15/9bc32671e9a38b413a76d24722a2bf8784a132c043063a8f5152d390b0f9/xxhash-3.6.0-cp313-cp313-win_amd64.whl", hash = "sha256:757320d45d2fbcce8f30c42a6b2f47862967aea7bf458b9625b4bbe7ee390392", size = 31542, upload-time = "2025-10-02T14:35:15.21Z" }, + { url = "https://files.pythonhosted.org/packages/39/c5/cc01e4f6188656e56112d6a8e0dfe298a16934b8c47a247236549a3f7695/xxhash-3.6.0-cp313-cp313-win_arm64.whl", hash = "sha256:457b8f85dec5825eed7b69c11ae86834a018b8e3df5e77783c999663da2f96d6", size = 27880, upload-time = "2025-10-02T14:35:16.315Z" }, + { url = "https://files.pythonhosted.org/packages/f3/30/25e5321c8732759e930c555176d37e24ab84365482d257c3b16362235212/xxhash-3.6.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:a42e633d75cdad6d625434e3468126c73f13f7584545a9cf34e883aa1710e702", size = 32956, upload-time = "2025-10-02T14:35:17.413Z" }, + { url = "https://files.pythonhosted.org/packages/9f/3c/0573299560d7d9f8ab1838f1efc021a280b5ae5ae2e849034ef3dee18810/xxhash-3.6.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:568a6d743219e717b07b4e03b0a828ce593833e498c3b64752e0f5df6bfe84db", size = 31072, upload-time = "2025-10-02T14:35:18.844Z" }, + { url = "https://files.pythonhosted.org/packages/7a/1c/52d83a06e417cd9d4137722693424885cc9878249beb3a7c829e74bf7ce9/xxhash-3.6.0-cp313-cp313t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:bec91b562d8012dae276af8025a55811b875baace6af510412a5e58e3121bc54", size = 196409, upload-time = "2025-10-02T14:35:20.31Z" }, + { url = "https://files.pythonhosted.org/packages/e3/8e/c6d158d12a79bbd0b878f8355432075fc82759e356ab5a111463422a239b/xxhash-3.6.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:78e7f2f4c521c30ad5e786fdd6bae89d47a32672a80195467b5de0480aa97b1f", size = 215736, upload-time = "2025-10-02T14:35:21.616Z" }, + { url = "https://files.pythonhosted.org/packages/bc/68/c4c80614716345d55071a396cf03d06e34b5f4917a467faf43083c995155/xxhash-3.6.0-cp313-cp313t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:3ed0df1b11a79856df5ffcab572cbd6b9627034c1c748c5566fa79df9048a7c5", size = 214833, upload-time = "2025-10-02T14:35:23.32Z" }, + { url = "https://files.pythonhosted.org/packages/7e/e9/ae27c8ffec8b953efa84c7c4a6c6802c263d587b9fc0d6e7cea64e08c3af/xxhash-3.6.0-cp313-cp313t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:0e4edbfc7d420925b0dd5e792478ed393d6e75ff8fc219a6546fb446b6a417b1", size = 448348, upload-time = "2025-10-02T14:35:25.111Z" }, + { url = "https://files.pythonhosted.org/packages/d7/6b/33e21afb1b5b3f46b74b6bd1913639066af218d704cc0941404ca717fc57/xxhash-3.6.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fba27a198363a7ef87f8c0f6b171ec36b674fe9053742c58dd7e3201c1ab30ee", size = 196070, upload-time = "2025-10-02T14:35:26.586Z" }, + { url = "https://files.pythonhosted.org/packages/96/b6/fcabd337bc5fa624e7203aa0fa7d0c49eed22f72e93229431752bddc83d9/xxhash-3.6.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:794fe9145fe60191c6532fa95063765529770edcdd67b3d537793e8004cabbfd", size = 212907, upload-time = "2025-10-02T14:35:28.087Z" }, + { url = "https://files.pythonhosted.org/packages/4b/d3/9ee6160e644d660fcf176c5825e61411c7f62648728f69c79ba237250143/xxhash-3.6.0-cp313-cp313t-musllinux_1_2_i686.whl", hash = "sha256:6105ef7e62b5ac73a837778efc331a591d8442f8ef5c7e102376506cb4ae2729", size = 200839, upload-time = "2025-10-02T14:35:29.857Z" }, + { url = "https://files.pythonhosted.org/packages/0d/98/e8de5baa5109394baf5118f5e72ab21a86387c4f89b0e77ef3e2f6b0327b/xxhash-3.6.0-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:f01375c0e55395b814a679b3eea205db7919ac2af213f4a6682e01220e5fe292", size = 213304, upload-time = "2025-10-02T14:35:31.222Z" }, + { url = "https://files.pythonhosted.org/packages/7b/1d/71056535dec5c3177eeb53e38e3d367dd1d16e024e63b1cee208d572a033/xxhash-3.6.0-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:d706dca2d24d834a4661619dcacf51a75c16d65985718d6a7d73c1eeeb903ddf", size = 416930, upload-time = "2025-10-02T14:35:32.517Z" }, + { url = "https://files.pythonhosted.org/packages/dc/6c/5cbde9de2cd967c322e651c65c543700b19e7ae3e0aae8ece3469bf9683d/xxhash-3.6.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:5f059d9faeacd49c0215d66f4056e1326c80503f51a1532ca336a385edadd033", size = 193787, upload-time = "2025-10-02T14:35:33.827Z" }, + { url = "https://files.pythonhosted.org/packages/19/fa/0172e350361d61febcea941b0cc541d6e6c8d65d153e85f850a7b256ff8a/xxhash-3.6.0-cp313-cp313t-win32.whl", hash = "sha256:1244460adc3a9be84731d72b8e80625788e5815b68da3da8b83f78115a40a7ec", size = 30916, upload-time = "2025-10-02T14:35:35.107Z" }, + { url = "https://files.pythonhosted.org/packages/ad/e6/e8cf858a2b19d6d45820f072eff1bea413910592ff17157cabc5f1227a16/xxhash-3.6.0-cp313-cp313t-win_amd64.whl", hash = "sha256:b1e420ef35c503869c4064f4a2f2b08ad6431ab7b229a05cce39d74268bca6b8", size = 31799, upload-time = "2025-10-02T14:35:36.165Z" }, + { url = "https://files.pythonhosted.org/packages/56/15/064b197e855bfb7b343210e82490ae672f8bc7cdf3ddb02e92f64304ee8a/xxhash-3.6.0-cp313-cp313t-win_arm64.whl", hash = "sha256:ec44b73a4220623235f67a996c862049f375df3b1052d9899f40a6382c32d746", size = 28044, upload-time = "2025-10-02T14:35:37.195Z" }, +] + +[[package]] +name = "yourdfpy" +version = "0.0.60" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "lxml" }, + { name = "numpy" }, + { name = "six" }, + { name = "trimesh", extra = ["easy"] }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ff/19/20c50861f30aff7720f9a601f386d73760c2df9961de1f98d0dbf3b85e69/yourdfpy-0.0.60.tar.gz", hash = "sha256:2af2d8bdeea1b85b642590a3b4236fdb35746d7b3e38ce460a169c18d9c4f868", size = 538238, upload-time = "2026-01-23T07:32:47.856Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/60/4ea0d6df0b497d51bf2ef87eaab0eb26f8bc3b3313c012da5df3383cced9/yourdfpy-0.0.60-py3-none-any.whl", hash = "sha256:8a187a8b18c98db87c76e9a950581b3c875b761e00df83942526c17ea693166c", size = 22194, upload-time = "2026-01-23T07:32:46.481Z" }, +] + +[[package]] +name = "zipp" +version = "3.23.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/30/21/093488dfc7cc8964ded15ab726fad40f25fd3d788fd741cc1c5a17d78ee8/zipp-3.23.1.tar.gz", hash = "sha256:32120e378d32cd9714ad503c1d024619063ec28aad2248dc6672ad13edfa5110", size = 25965, upload-time = "2026-04-13T23:21:46.6Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/08/8a/0861bec20485572fbddf3dfba2910e38fe249796cb73ecdeb74e07eeb8d3/zipp-3.23.1-py3-none-any.whl", hash = "sha256:0b3596c50a5c700c9cb40ba8d86d9f2cc4807e9bedb06bcdf7fac85633e444dc", size = 10378, upload-time = "2026-04-13T23:21:45.386Z" }, +] + +[[package]] +name = "zstandard" +version = "0.25.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/aa/3e0508d5a5dd96529cdc5a97011299056e14c6505b678fd58938792794b1/zstandard-0.25.0.tar.gz", hash = "sha256:7713e1179d162cf5c7906da876ec2ccb9c3a9dcbdffef0cc7f70c3667a205f0b", size = 711513, upload-time = "2025-09-14T22:15:54.002Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/fc/f26eb6ef91ae723a03e16eddb198abcfce2bc5a42e224d44cc8b6765e57e/zstandard-0.25.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7b3c3a3ab9daa3eed242d6ecceead93aebbb8f5f84318d82cee643e019c4b73b", size = 795738, upload-time = "2025-09-14T22:16:56.237Z" }, + { url = "https://files.pythonhosted.org/packages/aa/1c/d920d64b22f8dd028a8b90e2d756e431a5d86194caa78e3819c7bf53b4b3/zstandard-0.25.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:913cbd31a400febff93b564a23e17c3ed2d56c064006f54efec210d586171c00", size = 640436, upload-time = "2025-09-14T22:16:57.774Z" }, + { url = "https://files.pythonhosted.org/packages/53/6c/288c3f0bd9fcfe9ca41e2c2fbfd17b2097f6af57b62a81161941f09afa76/zstandard-0.25.0-cp312-cp312-manylinux2010_i686.manylinux2014_i686.manylinux_2_12_i686.manylinux_2_17_i686.whl", hash = "sha256:011d388c76b11a0c165374ce660ce2c8efa8e5d87f34996aa80f9c0816698b64", size = 5343019, upload-time = "2025-09-14T22:16:59.302Z" }, + { url = "https://files.pythonhosted.org/packages/1e/15/efef5a2f204a64bdb5571e6161d49f7ef0fffdbca953a615efbec045f60f/zstandard-0.25.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dffecc361d079bb48d7caef5d673c88c8988d3d33fb74ab95b7ee6da42652ea", size = 5063012, upload-time = "2025-09-14T22:17:01.156Z" }, + { url = "https://files.pythonhosted.org/packages/b7/37/a6ce629ffdb43959e92e87ebdaeebb5ac81c944b6a75c9c47e300f85abdf/zstandard-0.25.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:7149623bba7fdf7e7f24312953bcf73cae103db8cae49f8154dd1eadc8a29ecb", size = 5394148, upload-time = "2025-09-14T22:17:03.091Z" }, + { url = "https://files.pythonhosted.org/packages/e3/79/2bf870b3abeb5c070fe2d670a5a8d1057a8270f125ef7676d29ea900f496/zstandard-0.25.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:6a573a35693e03cf1d67799fd01b50ff578515a8aeadd4595d2a7fa9f3ec002a", size = 5451652, upload-time = "2025-09-14T22:17:04.979Z" }, + { url = "https://files.pythonhosted.org/packages/53/60/7be26e610767316c028a2cbedb9a3beabdbe33e2182c373f71a1c0b88f36/zstandard-0.25.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5a56ba0db2d244117ed744dfa8f6f5b366e14148e00de44723413b2f3938a902", size = 5546993, upload-time = "2025-09-14T22:17:06.781Z" }, + { url = "https://files.pythonhosted.org/packages/85/c7/3483ad9ff0662623f3648479b0380d2de5510abf00990468c286c6b04017/zstandard-0.25.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:10ef2a79ab8e2974e2075fb984e5b9806c64134810fac21576f0668e7ea19f8f", size = 5046806, upload-time = "2025-09-14T22:17:08.415Z" }, + { url = "https://files.pythonhosted.org/packages/08/b3/206883dd25b8d1591a1caa44b54c2aad84badccf2f1de9e2d60a446f9a25/zstandard-0.25.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:aaf21ba8fb76d102b696781bddaa0954b782536446083ae3fdaa6f16b25a1c4b", size = 5576659, upload-time = "2025-09-14T22:17:10.164Z" }, + { url = "https://files.pythonhosted.org/packages/9d/31/76c0779101453e6c117b0ff22565865c54f48f8bd807df2b00c2c404b8e0/zstandard-0.25.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:1869da9571d5e94a85a5e8d57e4e8807b175c9e4a6294e3b66fa4efb074d90f6", size = 4953933, upload-time = "2025-09-14T22:17:11.857Z" }, + { url = "https://files.pythonhosted.org/packages/18/e1/97680c664a1bf9a247a280a053d98e251424af51f1b196c6d52f117c9720/zstandard-0.25.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:809c5bcb2c67cd0ed81e9229d227d4ca28f82d0f778fc5fea624a9def3963f91", size = 5268008, upload-time = "2025-09-14T22:17:13.627Z" }, + { url = "https://files.pythonhosted.org/packages/1e/73/316e4010de585ac798e154e88fd81bb16afc5c5cb1a72eeb16dd37e8024a/zstandard-0.25.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:f27662e4f7dbf9f9c12391cb37b4c4c3cb90ffbd3b1fb9284dadbbb8935fa708", size = 5433517, upload-time = "2025-09-14T22:17:16.103Z" }, + { url = "https://files.pythonhosted.org/packages/5b/60/dd0f8cfa8129c5a0ce3ea6b7f70be5b33d2618013a161e1ff26c2b39787c/zstandard-0.25.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:99c0c846e6e61718715a3c9437ccc625de26593fea60189567f0118dc9db7512", size = 5814292, upload-time = "2025-09-14T22:17:17.827Z" }, + { url = "https://files.pythonhosted.org/packages/fc/5f/75aafd4b9d11b5407b641b8e41a57864097663699f23e9ad4dbb91dc6bfe/zstandard-0.25.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:474d2596a2dbc241a556e965fb76002c1ce655445e4e3bf38e5477d413165ffa", size = 5360237, upload-time = "2025-09-14T22:17:19.954Z" }, + { url = "https://files.pythonhosted.org/packages/ff/8d/0309daffea4fcac7981021dbf21cdb2e3427a9e76bafbcdbdf5392ff99a4/zstandard-0.25.0-cp312-cp312-win32.whl", hash = "sha256:23ebc8f17a03133b4426bcc04aabd68f8236eb78c3760f12783385171b0fd8bd", size = 436922, upload-time = "2025-09-14T22:17:24.398Z" }, + { url = "https://files.pythonhosted.org/packages/79/3b/fa54d9015f945330510cb5d0b0501e8253c127cca7ebe8ba46a965df18c5/zstandard-0.25.0-cp312-cp312-win_amd64.whl", hash = "sha256:ffef5a74088f1e09947aecf91011136665152e0b4b359c42be3373897fb39b01", size = 506276, upload-time = "2025-09-14T22:17:21.429Z" }, + { url = "https://files.pythonhosted.org/packages/ea/6b/8b51697e5319b1f9ac71087b0af9a40d8a6288ff8025c36486e0c12abcc4/zstandard-0.25.0-cp312-cp312-win_arm64.whl", hash = "sha256:181eb40e0b6a29b3cd2849f825e0fa34397f649170673d385f3598ae17cca2e9", size = 462679, upload-time = "2025-09-14T22:17:23.147Z" }, + { url = "https://files.pythonhosted.org/packages/35/0b/8df9c4ad06af91d39e94fa96cc010a24ac4ef1378d3efab9223cc8593d40/zstandard-0.25.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:ec996f12524f88e151c339688c3897194821d7f03081ab35d31d1e12ec975e94", size = 795735, upload-time = "2025-09-14T22:17:26.042Z" }, + { url = "https://files.pythonhosted.org/packages/3f/06/9ae96a3e5dcfd119377ba33d4c42a7d89da1efabd5cb3e366b156c45ff4d/zstandard-0.25.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:a1a4ae2dec3993a32247995bdfe367fc3266da832d82f8438c8570f989753de1", size = 640440, upload-time = "2025-09-14T22:17:27.366Z" }, + { url = "https://files.pythonhosted.org/packages/d9/14/933d27204c2bd404229c69f445862454dcc101cd69ef8c6068f15aaec12c/zstandard-0.25.0-cp313-cp313-manylinux2010_i686.manylinux2014_i686.manylinux_2_12_i686.manylinux_2_17_i686.whl", hash = "sha256:e96594a5537722fdfb79951672a2a63aec5ebfb823e7560586f7484819f2a08f", size = 5343070, upload-time = "2025-09-14T22:17:28.896Z" }, + { url = "https://files.pythonhosted.org/packages/6d/db/ddb11011826ed7db9d0e485d13df79b58586bfdec56e5c84a928a9a78c1c/zstandard-0.25.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:bfc4e20784722098822e3eee42b8e576b379ed72cca4a7cb856ae733e62192ea", size = 5063001, upload-time = "2025-09-14T22:17:31.044Z" }, + { url = "https://files.pythonhosted.org/packages/db/00/87466ea3f99599d02a5238498b87bf84a6348290c19571051839ca943777/zstandard-0.25.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:457ed498fc58cdc12fc48f7950e02740d4f7ae9493dd4ab2168a47c93c31298e", size = 5394120, upload-time = "2025-09-14T22:17:32.711Z" }, + { url = "https://files.pythonhosted.org/packages/2b/95/fc5531d9c618a679a20ff6c29e2b3ef1d1f4ad66c5e161ae6ff847d102a9/zstandard-0.25.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:fd7a5004eb1980d3cefe26b2685bcb0b17989901a70a1040d1ac86f1d898c551", size = 5451230, upload-time = "2025-09-14T22:17:34.41Z" }, + { url = "https://files.pythonhosted.org/packages/63/4b/e3678b4e776db00f9f7b2fe58e547e8928ef32727d7a1ff01dea010f3f13/zstandard-0.25.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8e735494da3db08694d26480f1493ad2cf86e99bdd53e8e9771b2752a5c0246a", size = 5547173, upload-time = "2025-09-14T22:17:36.084Z" }, + { url = "https://files.pythonhosted.org/packages/4e/d5/ba05ed95c6b8ec30bd468dfeab20589f2cf709b5c940483e31d991f2ca58/zstandard-0.25.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:3a39c94ad7866160a4a46d772e43311a743c316942037671beb264e395bdd611", size = 5046736, upload-time = "2025-09-14T22:17:37.891Z" }, + { url = "https://files.pythonhosted.org/packages/50/d5/870aa06b3a76c73eced65c044b92286a3c4e00554005ff51962deef28e28/zstandard-0.25.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:172de1f06947577d3a3005416977cce6168f2261284c02080e7ad0185faeced3", size = 5576368, upload-time = "2025-09-14T22:17:40.206Z" }, + { url = "https://files.pythonhosted.org/packages/5d/35/398dc2ffc89d304d59bc12f0fdd931b4ce455bddf7038a0a67733a25f550/zstandard-0.25.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:3c83b0188c852a47cd13ef3bf9209fb0a77fa5374958b8c53aaa699398c6bd7b", size = 4954022, upload-time = "2025-09-14T22:17:41.879Z" }, + { url = "https://files.pythonhosted.org/packages/9a/5c/36ba1e5507d56d2213202ec2b05e8541734af5f2ce378c5d1ceaf4d88dc4/zstandard-0.25.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:1673b7199bbe763365b81a4f3252b8e80f44c9e323fc42940dc8843bfeaf9851", size = 5267889, upload-time = "2025-09-14T22:17:43.577Z" }, + { url = "https://files.pythonhosted.org/packages/70/e8/2ec6b6fb7358b2ec0113ae202647ca7c0e9d15b61c005ae5225ad0995df5/zstandard-0.25.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:0be7622c37c183406f3dbf0cba104118eb16a4ea7359eeb5752f0794882fc250", size = 5433952, upload-time = "2025-09-14T22:17:45.271Z" }, + { url = "https://files.pythonhosted.org/packages/7b/01/b5f4d4dbc59ef193e870495c6f1275f5b2928e01ff5a81fecb22a06e22fb/zstandard-0.25.0-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:5f5e4c2a23ca271c218ac025bd7d635597048b366d6f31f420aaeb715239fc98", size = 5814054, upload-time = "2025-09-14T22:17:47.08Z" }, + { url = "https://files.pythonhosted.org/packages/b2/e5/fbd822d5c6f427cf158316d012c5a12f233473c2f9c5fe5ab1ae5d21f3d8/zstandard-0.25.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:4f187a0bb61b35119d1926aee039524d1f93aaf38a9916b8c4b78ac8514a0aaf", size = 5360113, upload-time = "2025-09-14T22:17:48.893Z" }, + { url = "https://files.pythonhosted.org/packages/8e/e0/69a553d2047f9a2c7347caa225bb3a63b6d7704ad74610cb7823baa08ed7/zstandard-0.25.0-cp313-cp313-win32.whl", hash = "sha256:7030defa83eef3e51ff26f0b7bfb229f0204b66fe18e04359ce3474ac33cbc09", size = 436936, upload-time = "2025-09-14T22:17:52.658Z" }, + { url = "https://files.pythonhosted.org/packages/d9/82/b9c06c870f3bd8767c201f1edbdf9e8dc34be5b0fbc5682c4f80fe948475/zstandard-0.25.0-cp313-cp313-win_amd64.whl", hash = "sha256:1f830a0dac88719af0ae43b8b2d6aef487d437036468ef3c2ea59c51f9d55fd5", size = 506232, upload-time = "2025-09-14T22:17:50.402Z" }, + { url = "https://files.pythonhosted.org/packages/d4/57/60c3c01243bb81d381c9916e2a6d9e149ab8627c0c7d7abb2d73384b3c0c/zstandard-0.25.0-cp313-cp313-win_arm64.whl", hash = "sha256:85304a43f4d513f5464ceb938aa02c1e78c2943b29f44a750b48b25ac999a049", size = 462671, upload-time = "2025-09-14T22:17:51.533Z" }, +]